when using crop with LazyImage the image data is not a tensor

romainVala · Nov 5, 2024 · 354dbf3 · 354dbf3
1 parent 2fe3eb2
commit 354dbf3
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/torchio/transforms/preprocessing/spatial/crop.py b/src/torchio/transforms/preprocessing/spatial/crop.py
@@ -1,5 +1,6 @@
 import nibabel as nib
 import numpy as np
+import torch
 
 from ....data.subject import Subject
 from .bounds_transform import BoundsTransform
@@ -48,7 +49,10 @@ def apply_transform(self, sample) -> Subject:
             new_affine[:3, 3] = new_origin
             i0, j0, k0 = index_ini
             i1, j1, k1 = index_fin
-            image.set_data(image.data[:, i0:i1, j0:j1, k0:k1].clone())
+            if isinstance(image.data,torch.Tensor):
+                image.set_data(image.data[:, i0:i1, j0:j1, k0:k1].clone())
+            else:
+                image.set_data(torch.as_tensor(image.data[:, i0:i1, j0:j1, k0:k1]))
             image.affine = new_affine
         return sample