fix: correct shape.

mind-inria · Oct 16, 2023 · e9b9885 · e9b9885
1 parent d7f1261
commit e9b9885
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/src/mrinufft/operators/stacked.py b/src/mrinufft/operators/stacked.py
@@ -372,7 +372,7 @@ def _op_sense_host(self, data, ksp):
         if ksp is None:
             ksp = np.empty((B, C, NZ * NS), dtype=self.cpx_dtype)
         ksp = ksp.reshape((B * C, NZ * NS))
-        ksp_batched = cp.empty((T, NZ * NS), dtype=self.cpx_dtype)
+        ksp_batched = cp.empty((T, NZ, NS), dtype=self.cpx_dtype)
         for i in range(B * C // T):
             idx_coils = np.arange(i * T, (i + 1) * T) % C
             idx_batch = np.arange(i * T, (i + 1) * T) // C
@@ -393,7 +393,6 @@ def _op_sense_host(self, data, ksp):
             # After reordering, apply 2D NUFFT
             ksp_batched = self.operator._op_calibless_device(cp.ascontiguousarray(tmp))
             ksp[i * T : (i + 1) * T] = ksp_batched.get()
-        ksp = ksp.reshape((B, C, NZ * NS))
         ksp = ksp.reshape((B, C, NZ, NS))
         return ksp