ADLR/megatron-lm!2416 - tests: Add barrier for destroy

NVIDIA · Dec 3, 2024 · 21cc9b0 · 21cc9b0
1 parent ae832c7
commit 21cc9b0
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tests/unit_tests/conftest.py b/tests/unit_tests/conftest.py
@@ -19,6 +19,7 @@ def pytest_sessionfinish(session, exitstatus):
 def cleanup():
     yield
     if torch.distributed.is_initialized():
+        torch.distributed.barrier()
         torch.distributed.destroy_process_group()