Added torch_ucc support (facebookresearch#173)

Adding Torch-UCC: https://github.com/openucx/torch-ucc import
jianan-gu · May 2, 2021 · e9b659e · e9b659e
1 parent 833ca3c
commit e9b659e
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/extend_distributed.py b/extend_distributed.py
@@ -20,6 +20,12 @@
     # print(e)
     torch_ccl = False
 
+try:
+    import torch_ucc
+except ImportError as e:
+    torch_ucc = False
+
+
 my_rank = -1
 my_size = -1
 my_local_rank = -1