condition if tpu_id is None

lezwon · Jun 1, 2020 · 1fb6e58 · 1fb6e58
1 parent 1741fc3
commit 1fb6e58
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/pytorch_lightning/trainer/evaluation_loop.py b/pytorch_lightning/trainer/evaluation_loop.py
@@ -250,7 +250,7 @@ def _evaluate(self, model: LightningModule, dataloaders, max_batches: int, test_
 
             # on TPU we have to wrap it under the ParallelLoader
             if self.use_tpu:
-                device = xm.xla_device(self.tpu_id)
+                device = xm.xla_device(self.tpu_id) if self.tpu_id is not None else xm.xla_device()
                 dataloader = xla_pl.ParallelLoader(dataloader, [device])
                 dataloader = dataloader.per_device_loader(device)
 

diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py
@@ -413,7 +413,7 @@ def run_training_epoch(self):
 
         # on TPU we have to wrap it under the ParallelLoader
         if self.use_tpu:
-            device = xm.xla_device(self.tpu_id)
+            device = xm.xla_device(self.tpu_id) if self.tpu_id is not None else xm.xla_device()
             train_dataloader = xla_pl.ParallelLoader(train_dataloader, [device])
             train_dataloader = train_dataloader.per_device_loader(device)