condition if tpu_id is None

Author: lezwon

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)
 

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)