This commit is contained in:
William Falcon 2020-04-23 15:03:39 -04:00 committed by GitHub
parent 29ebe92208
commit 68ca577919
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 7 additions and 2 deletions

5
.gitignore vendored
View File

@ -131,3 +131,8 @@ ENV/
.data/
datasets/
mnist/
# pl tests
ml-runs/
*.zip
pytorch\ lightning

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)