fix: fix train dataset tensor conversion

Author: okunator

cellseg_models_pytorch/torch_datasets/folder_dataset_train.py

@@ -6,7 +6,7 @@
 from torch.utils.data import Dataset
 
 from cellseg_models_pytorch.transforms.albu_transforms import ApplyEach
-from cellseg_models_pytorch.utils import FileHandler
+from cellseg_models_pytorch.utils import FileHandler, to_tensor
 
 try:
     import albumentations as A
@@ -118,8 +118,8 @@ def __getitem__(self, ix: int) -> Dict[str, np.ndarray]:
 
         tr = self.transforms(image=data["image"], masks=[masks])
 
-        image = tr["image"].to(self.output_device)
-        masks = tr["masks"][0].to(self.output_device)
+        image = to_tensor(tr["image"])
+        masks = to_tensor(tr["masks"][0])
         masks = torch.split(masks, mask_chls, dim=0)
 
         integer_masks = {k: masks[i] for i, k in enumerate(self.mask_keys)}

cellseg_models_pytorch/torch_datasets/hdf5_dataset_train.py

@@ -5,7 +5,7 @@
 from torch.utils.data import Dataset
 
 from cellseg_models_pytorch.transforms.albu_transforms import ApplyEach
-from cellseg_models_pytorch.utils import FileHandler
+from cellseg_models_pytorch.utils import FileHandler, to_tensor
 
 try:
     import albumentations as A
@@ -125,8 +125,8 @@ def __getitem__(self, ix: int) -> Dict[str, np.ndarray]:
 
         tr = self.transforms(image=data["image"], masks=[masks])
 
-        image = tr["image"].to(self.output_device)
-        masks = tr["masks"][0].to(self.output_device)
+        image = to_tensor(tr["image"])
+        masks = to_tensor(tr["masks"][0])
         masks = torch.split(masks, mask_chls, dim=0)
 
         integer_masks = {k: masks[i] for i, k in enumerate(self.mask_keys)}