Skip to content

Commit 6ca9257

Browse files
committed
Merge branch 'patch-1' of https://github.com/Delaunay/pytorch-image-models into Delaunay-dist-backend
2 parents 83aee5c + 8848dad commit 6ca9257

File tree

1 file changed

+7
-2
lines changed

1 file changed

+7
-2
lines changed

timm/utils/distributed.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -109,8 +109,13 @@ def init_distributed_device_so(
109109
global_rank = 0
110110
local_rank = 0
111111
if dist_backend is None:
112-
# FIXME sane defaults for other device backends?
113-
dist_backend = 'nccl' if 'cuda' in device else 'gloo'
112+
# FIXME: verify that ROCm transform nccl to rccl
113+
dist_backends = {
114+
"xpu": "ccl",
115+
"hpu": "hccl",
116+
"cuda": "nccl",
117+
}
118+
dist_backend = dist_backends.get(device, 'gloo')
114119
dist_url = dist_url or 'env://'
115120

116121
# TBD, support horovod?

0 commit comments

Comments
 (0)