Merge branch 'patch-1' of https://github.com/Delaunay/pytorch-image-m…

…odels into Delaunay-dist-backend
huggingface · May 15, 2024 · 6ca9257 · 6ca9257
2 parents 83aee5c + 8848dad
commit 6ca9257
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/timm/utils/distributed.py b/timm/utils/distributed.py
@@ -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?