fix: is_cuda_tensor guarded by TORCH_AVAILABLE.

mind-inria · Feb 5, 2025 · 48bc2ee · 48bc2ee
1 parent 9951fc6
commit 48bc2ee
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/src/mrinufft/_array_compat.py b/src/mrinufft/_array_compat.py
@@ -246,9 +246,9 @@ def _to_numpy_cupy(args, kwargs, leading_argument):
     This avoid transfers between different devices
     (e.g., CPU->GPU, GPU->CPU or different GPUs).
     """
-    if is_cuda_array(leading_argument) and CUPY_AVAILABLE:
-        return _to_cupy(*args, **kwargs)
-    elif is_cuda_tensor(leading_argument) and CUPY_AVAILABLE:
+    if (
+        is_cuda_array(leading_argument) or is_cuda_tensor(leading_argument)
+    ) and CUPY_AVAILABLE:
         return _to_cupy(*args, **kwargs)
     else:
         return _to_numpy(*args, **kwargs)

diff --git a/src/mrinufft/operators/interfaces/utils/gpu_utils.py b/src/mrinufft/operators/interfaces/utils/gpu_utils.py
@@ -36,7 +36,7 @@ def is_cuda_array(var):
 
 def is_cuda_tensor(var):
     """Check if var is a CUDA tensor."""
-    return isinstance(var, torch.Tensor) and var.is_cuda
+    return TORCH_AVAILABLE and isinstance(var, torch.Tensor) and var.is_cuda
 
 
 def is_host_array(var):