[r2] fix: move resnet_dt checking from graph to tabulate

Fix deepmodeling#3950. Backport a part of deepmodeling#3263 to r2 (the whole of deepmodeling#3263 is not likely to be backported). Signed-off-by: Jinzhe Zeng <[email protected]>
njzjz · Jul 5, 2024 · 897376a · 897376a
1 parent 84ca63c
commit 897376a
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/deepmd/utils/graph.py b/deepmd/utils/graph.py
@@ -176,10 +176,6 @@ def get_embedding_net_nodes_from_graph_def(
     embedding_net_nodes = get_pattern_nodes_from_graph_def(
         graph_def, embedding_net_pattern
     )
-    for key in embedding_net_nodes.keys():
-        assert (
-            key.find("bias") > 0 or key.find("matrix") > 0
-        ), "currently, only support weight matrix and bias matrix at the tabulation op!"
     return embedding_net_nodes
 
 

diff --git a/deepmd/utils/tabulate.py b/deepmd/utils/tabulate.py
@@ -133,6 +133,10 @@ def __init__(
         self.embedding_net_nodes = get_embedding_net_nodes_from_graph_def(
             self.graph_def, suffix=self.suffix
         )
+        for key in self.embedding_net_nodes.keys():
+            assert (
+                key.find("bias") > 0 or key.find("matrix") > 0
+            ), "currently, only support weight matrix and bias matrix at the tabulation op!"
 
         # move it to the descriptor class
         # for tt in self.exclude_types: