[SPMD] Fix TPU CI after #5478 (#5487)

* [SPMD] Fix TPU CI after #5478 Summary: Let's fix all TPU CI failures after #5478. Test Plan: TPU CI * Fix linters
pytorch · Aug 23, 2023 · 8a98a88 · 8a98a88
1 parent a2fab36
commit 8a98a88
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/test/pjrt/test_runtime_tpu.py b/test/pjrt/test_runtime_tpu.py
@@ -205,9 +205,11 @@ def test_global_runtime_device_attributes(self):
     results = pjrt.run_multiprocess(self._global_runtime_device_attributes)
     for result in results.values():
       for device in result:
-        self.assertCountEqual(['coords', 'core_on_chip'], list(device.keys()))
+        self.assertCountEqual(['coords', 'core_on_chip', 'name'],
+                              list(device.keys()))
         self.assertIsInstance(device['coords'], list)
         self.assertIsInstance(device['core_on_chip'], int)
+        self.assertIsInstance(device['name'], str)
 
   @staticmethod
   def _execute_time_metric():