mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Merge branch 'patch-1' of https://github.com/Delaunay/pytorch-image-models into Delaunay-dist-backend

Browse Source
This commit is contained in:
Ross Wightman 2024-05-15 08:40:58 -07:00
commit 6ca92570f7
1 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
timm/utils/distributed.py
View File

@ -109,8 +109,13 @@ def init_distributed_device_so(
global_rank = 0
local_rank = 0
if dist_backend is None:
# FIXME sane defaults for other device backends?
dist_backend = 'nccl' if 'cuda' in device else 'gloo'
# FIXME: verify that ROCm transform nccl to rccl
dist_backends = {
"xpu": "ccl",
"hpu": "hccl",
"cuda": "nccl",
}
dist_backend = dist_backends.get(device, 'gloo')
dist_url = dist_url or 'env://'
# TBD, support horovod?