Skip to content

GitLab

"src/vscode:/vscode.git/clone" did not exist on "37a787a106ff48e16ba64fa2df0c8fb446a6056e"
Unverified Commit 14e8e490 authored by wangjiangben-hw's avatar wangjiangben-hw Committed by GitHub
Browse files

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

parent b5134190
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
mmcv/runner/dist_utils.py
View file @ 14e8e490
...@@ -60,7 +60,8 @@ def _init_dist_pytorch(backend: str, **kwargs) -> None: ...@@ -60,7 +60,8 @@ def _init_dist_pytorch(backend: str, **kwargs) -> None:
**kwargs) **kwargs)
elif IS_NPU_AVAILABLE: elif IS_NPU_AVAILABLE:
import torch_npu # noqa: F401 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( dist.init_process_group(
backend='hccl', backend='hccl',
rank=rank, rank=rank,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment