still need to do intersection

Signed-off-by: Jinzhe Zeng <[email protected]>
deepmodeling · Mar 1, 2024 · 1fd9a7b · 1fd9a7b
1 parent a09b69d
commit 1fd9a7b
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/deepmd/tf/entrypoints/freeze.py b/deepmd/tf/entrypoints/freeze.py
@@ -365,13 +365,14 @@ def freeze_graph(
         "fitting_attr/aparam_nall",
         "spin_attr/ntypes_spin",
     ]
-    different_set = set(output_node) - set(input_node) - set(optional_node)
+    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
-        )
+        if different_set - set(optional_node):
+            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}")