Merge branch 'simple-cache' into gpu-transform

mehta-lab · Oct 31, 2024 · 70fcf1c · 70fcf1c
2 parents c185377 + 2ca134b
commit 70fcf1c
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/viscy/data/hcs_ram.py b/viscy/data/hcs_ram.py
@@ -318,7 +318,11 @@ def _setup_fit(self, dataset_settings: dict) -> None:
         )
 
     def train_dataloader(self) -> DataLoader:
-        sampler = ShardedDistributedSampler(self.train_dataset, shuffle=True)
+        if is_ddp_enabled():
+            sampler = ShardedDistributedSampler(self.train_dataset, shuffle=True)
+        else:
+            sampler = None
+            _logger.info("Using standard sampler for non-distributed training")
         return DataLoader(
             self.train_dataset,
             batch_size=self.batch_size // self.train_patches_per_stack,
@@ -333,7 +337,12 @@ def train_dataloader(self) -> DataLoader:
         )
 
     def val_dataloader(self) -> DataLoader:
-        sampler = ShardedDistributedSampler(self.val_dataset, shuffle=False)
+        if is_ddp_enabled():
+            sampler = ShardedDistributedSampler(self.val_dataset, shuffle=False)
+        else:
+            sampler = None
+            _logger.info("Using standard sampler for non-distributed validation")
+
         return DataLoader(
             self.val_dataset,
             batch_size=self.batch_size,