only freeze in rank 0 (#2937)

Fix #2272. Signed-off-by: Jinzhe Zeng <[email protected]>
deepmodeling · Oct 19, 2023 · b588928 · b588928
1 parent 64a5aa8
commit b588928
Showing 1 changed file with 5 additions and 1 deletion.
diff --git 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
     )