mirrors/mmengine
1
0
Fork 0
mirror of https://github.com/open-mmlab/mmengine.git synced 2025-06-03 21:54:44 +08:00
Code Issues Projects Releases Wiki Activity

[Fix] Fix multi card issue in PyTorch v2.1 on Ascend (#1321)

Browse Source
This commit is contained in:
LRJKD 2023-08-25 10:35:58 +08:00 committed by GitHub
parent e1c6079d73
commit a53c2802a6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mmengine/dist/dist.py vendored
View File

@ -416,7 +416,7 @@ def _broadcast_object_list(object_list: List[Any],
is_hccl_backend = group_backend == 'hccl'
is_cncl_backend = group_backend == 'cncl'
if is_hccl_backend:
current_device = torch.npu.current_device()
current_device = torch.device('npu', torch.npu.current_device())
object_sizes_tensor = object_sizes_tensor.to(current_device)
elif is_cncl_backend:
current_device = torch.device('mlu', torch.mlu.current_device())