Create distribute.py

open-mmlab · Dec 25, 2023 · 26f427e · 26f427e
1 parent ab78255
commit 26f427e
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/mmdeploy/pytorch/functions/distribute.py b/mmdeploy/pytorch/functions/distribute.py
@@ -0,0 +1,12 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from typing import Optional
+import torch
+
+from mmdeploy.core import FUNCTION_REWRITER
+
+@FUNCTION_REWRITER.register_rewriter(func_name='torch.distributed')
+def distributed_rewriter():
+    """rewrite torch.distributed to support some embedding device for higher PyTorch"""
+    # check torch.distributed is available?
+    if not torch.distributed.is_available():
+        torch.distributed.ReduceOp = lambda: None