added TPU_AVAILABLE variable

Lightning-AI · Sep 6, 2020 · dabcf7f · dabcf7f
1 parent 628657b
commit dabcf7f
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/pytorch_lightning/utilities/xla_device_utils.py b/pytorch_lightning/utilities/xla_device_utils.py
@@ -3,6 +3,7 @@
 from multiprocessing import Process, Queue
 
 TORCHXLA_AVAILABLE = importlib.util.find_spec("torch_xla") is not None
+TPU_AVAILABLE = None
 if TORCHXLA_AVAILABLE:
     import torch_xla.core.xla_model as xm
     import torch_xla.distributed.xla_multiprocessing as xmp
@@ -46,4 +47,7 @@ def is_device_tpu():
 
 
 def tpu_device_exists():
-    return pl_multi_process(is_device_tpu)()
+    global TPU_AVAILABLE
+    if TPU_AVAILABLE is None:
+        TPU_AVAILABLE = pl_multi_process(is_device_tpu)()
+    return TPU_AVAILABLE