Make TPU detection more robust (#5271)

pytorch · Jul 10, 2023 · cba267f · cba267f
1 parent bfb85fa
commit cba267f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torch_xla/runtime.py b/torch_xla/runtime.py
@@ -42,7 +42,7 @@ def _maybe_select_default_device():
   logging.warning('For more information about the status of PJRT, see '
                   'https://github.com/pytorch/xla/blob/master/docs/xr.md')
   # Check for libtpu _and_ the TPU device
-  if torch_xla._found_libtpu and os.path.exists('/dev/accel0'):
+  if torch_xla._found_libtpu and tpu.num_available_chips() > 0:
     logging.warning('libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU.')
     os.environ[xenv.PJRT_DEVICE] = 'TPU'
   else: