[Fix] Fix set_device bug when using multi-machine multi-device (#2370)

14e8e490 · wangjiangben-hw · GitHub · b5134190 · 14e8e490
Unverified Commit 14e8e490 authored Oct 31, 2022 by wangjiangben-hw Committed by GitHub Oct 31, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

mmcv/runner/dist_utils.py mmcv/runner/dist_utils.py +2 -1

No files found.
--- a/mmcv/runner/dist_utils.py
+++ b/mmcv/runner/dist_utils.py
@@ -60,7 +60,8 @@ def _init_dist_pytorch(backend: str, **kwargs) -> None:
            **kwargs)
    elif IS_NPU_AVAILABLE:
        import torch_npu  # noqa: F401
-        torch.npu.set_device(rank)
+        num_npus = torch.npu.device_count()
+        torch.npu.set_device(rank % num_npus)
        dist.init_process_group(
            backend='hccl',
            rank=rank,