change test_ddp src code

717f02f3 · zhanggzh · 5b602850 · 717f02f3
Commit 717f02f3 authored May 27, 2025 by zhanggzh
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 0 deletions

tests/test_ddp.py tests/test_ddp.py +1 -0

No files found.
--- a/tests/test_ddp.py
+++ b/tests/test_ddp.py
@@ -44,6 +44,7 @@ def _run_distributed(func, world_size, args: Dict, script=__file__, env=dict()):

    for i in range(world_size):
        env["OMPI_COMM_WORLD_RANK"] = str(i)
+        env["CUDA_VISIBLE_DEVICES"] = str(i % device_count)
        p = subprocess.Popen(
            [sys.executable, script, func, json.dumps(args)],
            stdout=subprocess.PIPE,