Update xla.py

Author: dimzhead

src/lightning/fabric/accelerators/xla.py

@@ -107,7 +107,7 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No
 
 def _using_pjrt() -> bool:
     # `using_pjrt` is removed in torch_xla 2.5
-    if _XLA_GREATER_EQUAL_2_5:
+    if True:
         from torch_xla import runtime as xr
 
         return xr.device_type() is not None
@@ -117,6 +117,10 @@ def _using_pjrt() -> bool:
 
         return xr.using_pjrt()
 
+    from torch_xla.experimental import pjrt
+
+    return pjrt.using_pjrt()
+
 
 def _parse_tpu_devices(devices: Union[int, str, list[int]]) -> Union[int, list[int]]:
     """Parses the TPU devices given in the format as accepted by the