From 919a26fe41961402a571f67cfb066af3c57d05e6 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Mon, 17 Feb 2020 17:46:46 -0500 Subject: [PATCH] fix tpu transfer bug --- pytorch_lightning/trainer/distrib_parts.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pytorch_lightning/trainer/distrib_parts.py b/pytorch_lightning/trainer/distrib_parts.py index 8cec776a..7c476708 100644 --- a/pytorch_lightning/trainer/distrib_parts.py +++ b/pytorch_lightning/trainer/distrib_parts.py @@ -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