Remove bymistake add

mind-inria · Apr 11, 2024 · 65c6bf7 · 65c6bf7
1 parent 33ff82a
commit 65c6bf7
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/src/mrinufft/io/nsp.py b/src/mrinufft/io/nsp.py
@@ -253,7 +253,7 @@ def read_trajectory(
     grad_filename: str,
     dwell_time: float = DEFAULT_RASTER_TIME,
     num_adc_samples: int = None,
-    gamma: Gammas |  = Gammas.HYDROGEN,
+    gamma: Gammas | float = Gammas.HYDROGEN,
     raster_time: float = DEFAULT_RASTER_TIME,
     read_shots: bool = False,
     normalize_factor: float = KMAX,

diff --git a/src/mrinufft/operators/base.py b/src/mrinufft/operators/base.py
@@ -130,9 +130,11 @@ def wrapper(self, data, output=None, *args, **kwargs):
         if xp.__name__ == "torch" and is_cuda_array(data):
             # Move them to cupy
             data_ = cp.from_dlpack(data)
+            output_ = cp.from_dlpack(output) if output is not None else None
         elif xp.__name__ == "torch":
             # Move to numpy
             data_ = data.to("cpu").numpy()
+            output_ = output.to("cpu").numpy() if output is not None else None
         else:
             data_ = data
             output_ = output