From 3a0223d171dbb7a7980d00563487451b1236ea59 Mon Sep 17 00:00:00 2001 From: "Dr. Kashif Rasul" Date: Tue, 16 Mar 2021 11:09:34 +0100 Subject: [PATCH] make pin_memory True by default --- pts/model/estimator.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pts/model/estimator.py b/pts/model/estimator.py index a10429c..85863b9 100644 --- a/pts/model/estimator.py +++ b/pts/model/estimator.py @@ -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, )