make pin_memory True by default

2026-06-27 18:06:19 +08:00 · 2021-03-16 11:09:34 +01:00
parent c4f6445d11
commit 3a0223d171
1 changed files with 2 additions and 0 deletions
@@ -117,6 +117,7 @@ class PyTorchEstimator(Estimator):
            batch_size=self.trainer.batch_size,
            num_workers=num_workers,
            prefetch_factor=prefetch_factor,
+            pin_memory=True,
            **kwargs,
        )

@@ -136,6 +137,7 @@ class PyTorchEstimator(Estimator):
                batch_size=self.trainer.batch_size,
                num_workers=num_workers,
                prefetch_factor=prefetch_factor,
+                pin_memory=True,
                **kwargs,
            )