slow tpu train

CHANGELOG.md

@@ -20,6 +20,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Attribute `best_model_path` to `ModelCheckpoint` for storing and later retrieving the path to the best saved model file ([#1799](https://github.com/PyTorchLightning/pytorch-lightning/pull/1799)) 
 
+- Speed up single-core TPU training by loading data using `ParallelLoader` ([#2033](https://github.com/PyTorchLightning/pytorch-lightning/pull/2033))
+
 ### Changed
 
 - Allow user to select individual TPU core to train on ([#1729](https://github.com/PyTorchLightning/pytorch-lightning/pull/1729))

pytorch_lightning/trainer/evaluation_loop.py

@@ -249,8 +249,8 @@ def _evaluate(self, model: LightningModule, dataloaders, max_batches: int, test_
             dl_outputs = []
 
             # 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)
                 dataloader = xla_pl.ParallelLoader(dataloader, [device])
                 dataloader = dataloader.per_device_loader(device)
 

pytorch_lightning/trainer/training_loop.py

@@ -412,8 +412,8 @@ def run_training_epoch(self):
         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)