Merge pull request #322 from sef43/zero_workers

Disable persistent_workers with num_workers=0 and disable pin_memory in DataModule
torchmd · May 16, 2024 · d8813f5 · d8813f5
2 parents 75c462a + fe3036a
commit d8813f5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchmdnet/data.py b/torchmdnet/data.py
@@ -133,8 +133,8 @@ def _get_dataloader(self, dataset, stage, store_dataloader=True):
             dataset=dataset,
             batch_size=batch_size,
             num_workers=self.hparams["num_workers"],
-            persistent_workers=True,
-            pin_memory=True,
+            persistent_workers=True if self.hparams["num_workers"] > 0 else False,
+            pin_memory=False,
             shuffle=shuffle,
         )