From 26f427e1c303369fc40683d04771548b39a8f3b8 Mon Sep 17 00:00:00 2001
From: yinfan98 <32722923+yinfan98@users.noreply.github.com>
Date: Mon, 25 Dec 2023 16:58:50 +0800
Subject: [PATCH] Create distribute.py

---
 mmdeploy/pytorch/functions/distribute.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)
 create mode 100644 mmdeploy/pytorch/functions/distribute.py

diff --git a/mmdeploy/pytorch/functions/distribute.py b/mmdeploy/pytorch/functions/distribute.py
new file mode 100644
index 0000000000..a069d175f7
--- /dev/null
+++ 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