From 41b6cbb3ca8a3a43e091d7f0de4d0184a8870d19 Mon Sep 17 00:00:00 2001 From: karlinjf Date: Thu, 23 Apr 2020 14:28:20 -0400 Subject: [PATCH] Don't copy the batch when training on a single gpu (#1576) * fix * whitespace Co-authored-by: Josh Karlin --- pytorch_lightning/trainer/training_loop.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py index 5b3d13c7..c5fb1b97 100644 --- a/pytorch_lightning/trainer/training_loop.py +++ b/pytorch_lightning/trainer/training_loop.py @@ -754,7 +754,11 @@ class TrainerTrainLoopMixin(ABC): gpu_id = 0 if isinstance(self.data_parallel_device_ids, list): gpu_id = self.data_parallel_device_ids[0] - batch = self.transfer_batch_to_gpu(copy.copy(batch), gpu_id) + + # Don't copy the batch since there is a single gpu that the batch could + # be referenced from and if there are multiple optimizers the batch will + # wind up copying it to the same device repeatedly. + batch = self.transfer_batch_to_gpu(batch, gpu_id) args[0] = batch output = self.model.training_step(*args)