From b90158bb2ac39d310e6c99b02cf2bf04a5e4d0d1 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Wed, 18 Oct 2023 18:33:56 -0400
Subject: [PATCH] only freeze in rank 0

Fix #2272.

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
---
 deepmd/entrypoints/freeze.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/deepmd/entrypoints/freeze.py b/deepmd/entrypoints/freeze.py
index 11e0d55645..22f3cb80b4 100755
--- a/deepmd/entrypoints/freeze.py
+++ b/deepmd/entrypoints/freeze.py
@@ -511,9 +511,13 @@ def freeze(
     # We import the meta graph and retrieve a Saver
     try:
         # In case paralle training
-        import horovod.tensorflow as _  # noqa: F401
+        import horovod.tensorflow as HVD
     except ImportError:
         pass
+    else:
+        HVD.init()
+        if HVD.rank() > 0:
+            return
     saver = tf.train.import_meta_graph(
         f"{input_checkpoint}.meta", clear_devices=clear_devices
     )