From 3562aa5aae7009fea6e5b6b972230ca8c3ad7d16 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Mon, 17 Feb 2020 17:47:16 -0500 Subject: [PATCH] fix tpu transfer bug 2 --- pytorch_lightning/trainer/distrib_parts.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytorch_lightning/trainer/distrib_parts.py b/pytorch_lightning/trainer/distrib_parts.py index 7c476708..29bc8178 100644 --- a/pytorch_lightning/trainer/distrib_parts.py +++ b/pytorch_lightning/trainer/distrib_parts.py @@ -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