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

fix tpu transfer bug 2

Browse Source
This commit is contained in:
William Falcon
2020-02-17 17:47:16 -05:00
parent 919a26fe41
commit 3562aa5aae
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/trainer/distrib_parts.py
+1 -1
View File
@@ -449,7 +449,7 @@ class TrainerDPMixin(ABC):
# when dict
if isinstance(batch, dict):
for k, v in batch.items():
batch[k] = self.__transfer_data_to_device(x, device, gpu_id)
batch[k] = self.__transfer_data_to_device(v, device, gpu_id)
return batch