Merge pull request #124 from iKernels/hotfix-#123

fixed #123
iKernels · May 19, 2022 · 4ffb91b · 4ffb91b
2 parents 4e50f75 + 07a3b5a
commit 4ffb91b
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/tests/helpers.py b/tests/helpers.py
@@ -28,6 +28,7 @@
     weight_decay=0.1,
     padding='max_length',
     max_length=128,
+    pin_memory=False,
 )
 
 

diff --git a/transformers_lightning/datamodules/super_datamodule.py b/transformers_lightning/datamodules/super_datamodule.py
@@ -76,7 +76,7 @@ def default_dataloader(self, dataset: Dataset, batch_size: int, **kwargs):
             dataset,
             batch_size=batch_size,
             num_workers=self.hyperparameters.num_workers,
-            pin_memory=True,
+            pin_memory=self.hyperparameters.pin_memory,
             collate_fn=self.collate_fn,
             **kwargs,
         )
@@ -121,6 +121,7 @@ def add_datamodule_specific_args(parser: ArgumentParser):
             type=int,
             help='Number of workers to be used to load datasets'
         )
+        parser.add_argument('--pin_memory', action="store_true", help='Whether to use memory pinning.')
         parser.add_argument('--batch_size', type=int, default=32)
         parser.add_argument('--val_batch_size', type=int, default=256)
         parser.add_argument('--test_batch_size', type=int, default=256)