fix missed add elimination

inisis · Oct 31, 2024 · 8b1c808 · 8b1c808
1 parent b76595b
commit 8b1c808
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/onnxslim/core/optimization/dead_node_elimination.py b/onnxslim/core/optimization/dead_node_elimination.py
@@ -73,7 +73,7 @@ def dead_node_elimination(graph, is_subgraph=False):
                 if value.ndim == 0 and value == 0:
                     delete_node(node, var_idx)
                     logger.debug(f"removing {node.op} op: {node.name}")
-                elif np.all(value == 0) and (node.inputs[0].shape == node.outputs[0].shape):
+                elif np.all(value == 0) and (node.inputs[var_idx].shape == node.outputs[0].shape):
                     delete_node(node, var_idx)
                     logger.debug(f"removing {node.op} op: {node.name}")
         elif node.op == "Expand":