This commit is contained in:
William Falcon authored and GitHub committed 2020-04-23 15:03:39 -04:00
1 parent 29ebe92208
commit 68ca577919
2 files changed
+7 -2

No files matched your search

+5
View File
@@ -131,3 +131,8 @@ ENV/
.data/
datasets/
mnist/
# pl tests
ml-runs/
*.zip
pytorch\ lightning
+2 -2
View File
@@ -752,7 +752,7 @@ class TrainerTrainLoopMixin(ABC):
# Horovod
elif self.use_horovod and self.on_gpu:
batch = self.transfer_batch_to_gpu(copy.copy(batch), hvd.local_rank())
batch = self.transfer_batch_to_gpu(batch, hvd.local_rank())
args[0] = batch
output = self.model.training_step(*args)
@@ -771,7 +771,7 @@ class TrainerTrainLoopMixin(ABC):
# TPU support
elif self.use_tpu:
batch = self.transfer_batch_to_tpu(copy.copy(batch))
batch = self.transfer_batch_to_tpu(batch)
args[0] = batch
output = self.model.training_step(*args)