Merge pull request #28 from mehta-lab/target-center-slice-fix

Fix datamodule

viscy/light/data.py

@@ -450,7 +450,11 @@ def _setup_predict(self, dataset_settings: dict):
         )
 
     def on_before_batch_transfer(self, batch: Sample, dataloader_idx: int) -> Sample:
-        if self.trainer.predicting or isinstance(batch, torch.Tensor):
+        predicting = False
+        if self.trainer:
+            if self.trainer.predicting:
+                predicting = True
+        if predicting or isinstance(batch, torch.Tensor):
             # skipping example input array
             return batch
         if self.target_2d:
@@ -465,7 +469,7 @@ def train_dataloader(self):
             batch_size=self.batch_size,
             num_workers=self.num_workers,
             shuffle=True,
-            persistent_workers=True,
+            persistent_workers=bool(self.num_workers),
         )
 
     def val_dataloader(self):
@@ -474,7 +478,7 @@ def val_dataloader(self):
             batch_size=self.batch_size,
             num_workers=self.num_workers,
             shuffle=False,
-            persistent_workers=True,
+            persistent_workers=bool(self.num_workers),
         )
 
     def test_dataloader(self):