wassname/pytorch-lightning
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-lightning.git synced 2026-09-11 12:31:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

Don't copy the batch when training on a single gpu (#1576)

Browse Source 
* fix

* whitespace

Co-authored-by: Josh Karlin <karlinjf@gmail.com>
This commit is contained in:
karlinjf
2020-04-23 14:28:20 -04:00
committed by GitHub
co-authored by Josh Karlin
parent 0b22b64a10
commit 41b6cbb3ca
1 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/trainer/training_loop.py
+5 -1
View File
@@ -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)