Ensure dataset is loaded onto CPU when pin_memory is used

Author: rhoadesScholar

src/cellmap_data/dataloader.py

@@ -231,6 +231,10 @@ def refresh(self):
 
         dataloader_kwargs.pop("force_has_data", None)
 
+        # Ensure that dataset is loaded onto CPU if pin_memory is used
+        if self._pin_memory:
+            self.dataset.to("cpu")
+
         self._pytorch_loader = torch.utils.data.DataLoader(
             self.dataset, **dataloader_kwargs
         )

src/cellmap_data/empty_image.py

@@ -81,7 +81,7 @@ def bg_count(self) -> float:
     def class_counts(self) -> float:
         return self._class_counts
 
-    def to(self, device: str, non_blocking: bool = True) -> None:
+    def to(self, device: str | torch.device, non_blocking: bool = True) -> None:
         self.store = self.store.to(device, non_blocking=non_blocking)
 
     def set_spatial_transforms(self, transforms: Mapping[str, Any] | None) -> None: