Revert "Disable persistent_workers with num_workers=0 and disable pin…

…_memory in DataModule "
torchmd · May 16, 2024 · f30743c · f30743c
1 parent d8813f5
commit f30743c
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 if self.hparams["num_workers"] > 0 else False,
-            pin_memory=False,
+            persistent_workers=True,
+            pin_memory=True,
             shuffle=shuffle,
         )