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

fix tpu transfer bug

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