remove xla_device call

Lightning-AI · Jun 7, 2020 · ca33af1 · ca33af1
1 parent a9114f8
commit ca33af1
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/tests/models/test_tpu.py b/tests/models/test_tpu.py
@@ -8,8 +8,7 @@
 
 try:
     import torch_xla
-    device = torch_xla.core.xla_model.xla_device()
-    device_type = torch_xla.core.xla_model.xla_device_hw(device)
+    device_type = torch_xla.core.xla_model.xla_device_hw('xla:1')
     TPU_AVAILABLE = device_type == 'TPU'
 except ImportError:
     TPU_AVAILABLE = False