tf: remove freeze warning for optional nodes (deepmodeling#3381)

Fix deepmodeling#3334. --------- Signed-off-by: Jinzhe Zeng <[email protected]> (cherry picked from commit 92ee632)
njzjz · Jul 2, 2024 · 644cc52 · 644cc52
1 parent f8a0b31
commit 644cc52
Showing 1 changed file with 13 additions and 5 deletions.
diff --git a/deepmd/entrypoints/freeze.py b/deepmd/entrypoints/freeze.py
@@ -357,13 +357,21 @@ def freeze_graph(
     output_node = _make_node_names(
         freeze_type, modifier, out_suffix=out_suffix, node_names=node_names
     )
+    # see #3334
+    optional_node = [
+        "train_attr/min_nbor_dist",
+        "fitting_attr/aparam_nall",
+        "spin_attr/ntypes_spin",
+    ]
     different_set = set(output_node) - set(input_node)
     if different_set:
-        log.warning(
-            "The following nodes are not in the graph: %s. "
-            "Skip freezeing these nodes. You may be freezing "
-            "a checkpoint generated by an old version." % different_set
-        )
+        different_set -= set(optional_node)
+        if different_set:
+            log.warning(
+                "The following nodes are not in the graph: %s. "
+                "Skip freezeing these nodes. You may be freezing "
+                "a checkpoint generated by an old version." % different_set
+            )
         # use intersection as output list
         output_node = list(set(output_node) & set(input_node))
     log.info(f"The following nodes will be frozen: {output_node}")