Add a condiction when we try to remove a transpose node. (#2272)

Signed-off-by: Jay Zhang <[email protected]>
onnx · Nov 28, 2023 · 07dfaa8 · 07dfaa8
1 parent ae4c39e
commit 07dfaa8
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/tf2onnx/optimizer/transpose_optimizer.py b/tf2onnx/optimizer/transpose_optimizer.py
@@ -509,10 +509,14 @@ def _add_handler(self, trans, node):
                 return True
         return self._handle_node_having_branches(trans, node)
 
+    def _output_node_has_single_consumer_node(self, node):
+        output_node = self._g.get_node_by_name(node.output[0])
+        return output_node and output_node.output and self._nodes_has_single_consumer_node([output_node])
+
     def _transpose_handler(self, trans, node):
         perm = trans.get_attr_value("perm")
         perm_inv = invert_perm(perm)
-        if is_tranpose_of_type(node, perm_inv):
+        if is_tranpose_of_type(node, perm_inv) and self._output_node_has_single_consumer_node(node):
             for g in {self._g, node.graph}:
                 g.replace_all_inputs(node.output[0], trans.input[0])  # ops=g.get_nodes()
 

diff --git a/tf2onnx/version.py b/tf2onnx/version.py
@@ -2,4 +2,4 @@
 
 
 version = '1.15.1'
-git_version = 'dc6155b52a137d858456fcc6bc720c327eec5612'
+git_version = 'ae4c39ed3bdab7edf487d73d5892a573684d1d6a'