Fix ClusterResolver

pytorch · Nov 22, 2023 · fd088d9 · fd088d9
1 parent f8ecbc2
commit fd088d9
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/benchmarks/util.py b/benchmarks/util.py
@@ -13,7 +13,7 @@
 
 try:
   import torch_xla.core.xla_model as xm
-  from torch_xla.distributed.cluster import ClusterResolver
+  from torch_xla._internal import tpu
 except ImportError:
   # ignore the error if torch_xla is not installed
   pass
@@ -150,5 +150,4 @@ def get_gpu_name():
 
 
 def get_tpu_name():
-  return ClusterResolver.get_instance_metadata(
-      'instance/attributes/accelerator-type')
+  return tpu._get_metadata("accelerator-type")