[Enhancement] Handle tensor device type in sync_random_seed (open-mml…

…ab#126)
HAOCHENYE · Mar 13, 2022 · f548c81 · f548c81
1 parent 6d73b6c
commit f548c81
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/mmengine/dist/dist.py b/mmengine/dist/dist.py
@@ -305,10 +305,16 @@ def sync_random_seed(group: Optional[dist.ProcessGroup] = None) -> int:
     if group is None:
         group = get_default_group()
 
+    group_backend = get_backend(group)
+    is_nccl_backend = group_backend == dist.Backend.NCCL
+    current_device = torch.device('cpu')
+    if is_nccl_backend:
+        current_device = torch.device('cuda', torch.cuda.current_device())
+
     if get_rank(group) == 0:
-        random_num = torch.tensor(seed, dtype=torch.int32)
+        random_num = torch.tensor(seed, dtype=torch.int32).to(current_device)
     else:
-        random_num = torch.tensor(0, dtype=torch.int32)
+        random_num = torch.tensor(0, dtype=torch.int32).to(current_device)
 
     dist.broadcast(random_num, src=0, group=group)
 

diff --git a/tests/test_dist/test_dist.py b/tests/test_dist/test_dist.py
@@ -190,8 +190,7 @@ def _test_broadcast_dist(device):
 
 def _test_sync_random_seed_dist(device):
     with patch.object(
-            torch, 'tensor',
-            return_value=torch.tensor(1024).to(device)) as mock_tensor:
+            torch, 'tensor', return_value=torch.tensor(1024)) as mock_tensor:
         output = dist.sync_random_seed()
         assert output == 1024
     mock_tensor.assert_called()