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

slow tpu train (#2033)

Browse Source 
* use parallel loader

* Revert "use parallel loader"

This reverts commit ed6e7583

* select tpu id for pl

* condition if tpu_id is None

* added info to changelog

* Revert "condition if tpu_id is None"

This reverts commit 1fb6e586

* Apply suggestions from code review

Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>
This commit is contained in:
Lezwon Castelino
2020-06-02 18:48:05 -04:00
committed by GitHub
co-authored by Jirka Borovec
parent fa696ce512
commit 943c4b20af
3 changed files with 6 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/trainer/training_loop.py
+2 -2
View File
@@ -415,8 +415,8 @@ class TrainerTrainLoopMixin(ABC):
train_dataloader = self.train_dataloader
# on TPU we have to wrap it under the ParallelLoader
if self.use_tpu and self.tpu_id is None:
device = xm.xla_device()
if self.use_tpu:
device = xm.xla_device(self.tpu_id)
train_dataloader = xla_pl.ParallelLoader(train_dataloader, [device])
train_dataloader = train_dataloader.per_device_loader(device)