Add missing import for XLA init_method (#6809)

pytorch · Mar 22, 2024 · 34d17c2 · 34d17c2
1 parent b0ceb2b
commit 34d17c2
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/test/spmd/test_xla_distributed_checkpoint.py b/test/spmd/test_xla_distributed_checkpoint.py
@@ -144,7 +144,8 @@ def test_resharding_different_device_mesh(self):
   def test_multihost_checkpoint(self):
     torch.manual_seed(42)
 
-    # Initialize the default CPU process group from the environment.
+    # Initialize the default CPU process group.
+    import torch_xla.distributed.xla_backend
     dist.init_process_group(backend='gloo', init_method='xla://')
 
     model1 = self._get_sharded_model(mesh_shape=(1, self.n_devices))