diff --git a/examples/dagrt-fusion.py b/examples/dagrt-fusion.py index 2a285ad533b0d1ffb545da65f69efc8665072a10..3cd6b71faf176db5a067b06d914ecd690909d10c 100755 --- a/examples/dagrt-fusion.py +++ b/examples/dagrt-fusion.py @@ -107,6 +107,8 @@ def dof_array_nbytes(ary: np.ndarray): return sum( dof_array_nbytes(ary[idx]) for idx in np.ndindex(ary.shape)) + elif isinstance(ary, DOFArray): + return sum(dof_array_nbytes(ary_i) for ary_i in ary) else: return ary.nbytes