change current_device position

DeepLink-org · Mar 30, 2024 · f510bd4 · f510bd4
1 parent f4e4d4a
commit f510bd4
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/dipu/tests/python/individual_scripts/test_rt_ddp.py b/dipu/tests/python/individual_scripts/test_rt_ddp.py
@@ -65,8 +65,9 @@ def forward(self, x):
 def demo_basic_ddp(rank, world_size, port):
     import torch_dipu
 
-    print(f"Running basic DDP example on rank {rank} {torch.cuda.current_device()}")
+    # print(f"Running basic DDP example on rank {rank} {torch.cuda.current_device()}")
     torch.cuda.set_device(rank)
+    print(f"Running basic DDP example on rank {rank} {torch.cuda.current_device()}")
     backend = "nccl"
     dev1 = rank
 
@@ -100,8 +101,9 @@ def demo_basic_ddp(rank, world_size, port):
 def demo_allreduce(rank, world_size, port):
     import torch_dipu
 
-    print(f"Running basic DDP example on rank {rank} {torch.cuda.current_device()}")
+    # print(f"Running basic DDP example on rank {rank} {torch.cuda.current_device()}")
     torch.cuda.set_device(rank)
+    print(f"Running basic DDP example on rank {rank} {torch.cuda.current_device()}")
     dev1 = rank
 
     setup(rank, world_size, port)
@@ -151,7 +153,7 @@ def demo_allgather(rank, world_size, port):
     import torch_dipu
 
     setup(rank, world_size, port)
-
+    print(f'rank={rank}')
     src1 = torch.ones((2, 4)).to(rank)
     dests = torch.zeros((world_size * 2, 4)).to(rank)
     dests = [