postprocess fix

bikefusion/data_utils.py

@@ -29,7 +29,7 @@ def pad_to_square(image, target_size=(128, 128)):
 
 def preprocess(images):
     # Convert arrays to tensors
-    images = torch.from_numpy(images).float()
+    images = torch.tensor(images).float()
 
     # Apply padding to each image in the dataset
     images = torch.stack([pad_to_square(img) for img in images])
@@ -52,7 +52,12 @@ def un_pad(image, target_size=(80, 128)):
 
 def postprocess(images):
     # Convert tensors to arrays
-    images = images.detach().cpu().numpy()
+    if isinstance(images, torch.Tensor):
+        images = images.detach().cpu().numpy()
+    elif isinstance(images, np.ndarray):
+        pass
+    else:
+        raise ValueError("images must be either a torch.Tensor or a np.ndarray")
 
     # Unpad each image in the dataset
     images = np.stack([un_pad(img) for img in images])