[WIP] fix tests

huggingface · May 13, 2024 · 107266b · 107266b
1 parent c7a8a86
commit 107266b
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/optimum/neuron/accelerate/state.py b/optimum/neuron/accelerate/state.py
@@ -85,7 +85,8 @@ def __init__(self, cpu: bool = False, **kwargs):
                 set_common_flags()
                 if os.environ.get("ACCELERATE_USE_AMP", "false") == "true":
                     set_neuron_cc_flags_for_torch_amp()
-                init_process_group()
+                if not torch.distributed.is_initialized():
+                    init_process_group()
                 self.num_processes = xm.xrt_world_size()
                 self.process_index = xm.get_ordinal()
                 self.local_process_index = xm.get_local_ordinal()