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

2022-10-31 22:56:48 +08:00 · 2022-10-31 22:56:48 +08:00 · 14e8e4904b
parent b51341909b
commit 14e8e4904b
1 changed files with 2 additions and 1 deletions
--- 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,