Fix internal CI failure

pytorch · Sep 22, 2022 · 149ce37 · 149ce37
1 parent 8900311
commit 149ce37
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/test/test_distributed.py b/test/test_distributed.py
@@ -6,7 +6,6 @@
 
 
 import os
-import subprocess
 import unittest
 
 from functools import partial
@@ -162,9 +161,11 @@ def test_distributed_dl2(self, backend) -> None:
     def test_elastic_training_dl2(self, backend) -> None:
         world_size = DEFAULT_WORLD_SIZE if backend != "nccl" else torch.cuda.device_count()
         nnodes = 1
-        subprocess.run(
+        from torch.distributed import run
+
+        run.main(
             [
-                "torchrun",
+                "--run_path",
                 f"--nnodes={nnodes}",
                 f"--nproc_per_node={world_size}",
                 abs_path("bin/elastic_training.py"),
@@ -187,9 +188,11 @@ def test_distributed_dl1(self, backend) -> None:
     def test_elastic_training_dl1(self, backend) -> None:
         world_size = DEFAULT_WORLD_SIZE if backend != "nccl" else torch.cuda.device_count()
         nnodes = 1
-        subprocess.run(
+        from torch.distributed import run
+
+        run.main(
             [
-                "torchrun",
+                "--run_path",
                 f"--nnodes={nnodes}",
                 f"--nproc_per_node={world_size}",
                 abs_path("bin/elastic_training.py"),