Fix errors

Author: ghpvnist

test/pytorch_test_base.py

@@ -559,7 +559,7 @@ def _alt_lookup(d, keys, defval):
   def instantiate_test(cls, name, test, *, generic_cls):
     test_name = name + '_' + cls.device_type
     class_name = cls.__name__
-    real_device_type = xm.xla_device_hw(str(torch.device('xla')))
+    real_device_type = xm.xla_device_hw(str(torch_xla.device()))
     assert real_device_type in DISABLED_TORCH_TESTS, 'Unsupported device type:' + real_device_type
     disabled_torch_tests = DISABLED_TORCH_TESTS[real_device_type]
 
@@ -632,7 +632,7 @@ def get_primary_device(cls):
   @classmethod
   def setUpClass(cls):
     # Sets the primary test device to the xla_device (CPU or TPU)
-    cls.primary_device = str(torch.device('xla'))
+    cls.primary_device = str(torch_xla.device())
     torch_xla._XLAC._xla_set_mat_mul_precision('highest')
 
   def setUp(self):

test/test_operations.py

@@ -440,12 +440,11 @@ def test_nonzero_cast(self):
 class TestOptimizationBarrier(test_utils.XlaTestCase):
 
   def test_optimization_barrier_correctness(self):
-    device = torch.device('xla')
     # only test optimization_barrier on TPU
-    if xm.xla_device_hw(device) != 'TPU':
+    if xm.xla_device_hw(torch_xla.device()) != 'TPU':
       return
-    x = torch.randn(5, 5, device=device)
-    y = torch.randn(5, 5, device=device)
+    x = torch.randn(5, 5, device='xla')
+    y = torch.randn(5, 5, device='xla')
     z = x + y
     xm.optimization_barrier_([x, y])
     self.assertEqual(z, x + y)

torch_xla/runtime.py

@@ -156,7 +156,8 @@ def local_ordinal() -> int:
   Local ordinal is in range [0, local_device_count)."""
   local_rank = xu.getenv_as(xenv.PJRT_LOCAL_PROCESS_RANK, int, 0)
   devices_per_process = addressable_device_count()
-  return local_rank * devices_per_process + torch.device('xla').index
+  return local_rank * devices_per_process + torch.device(
+      torch_xla._XLAC._xla_get_default_device()).index
 
 
 def process_index() -> int: