From 208cfa951e9e333daaef0d504b87df94b054a365 Mon Sep 17 00:00:00 2001 From: inisis Date: Tue, 3 Dec 2024 22:26:00 +0800 Subject: [PATCH] fix delete node bug --- onnxslim/core/utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/onnxslim/core/utils.py b/onnxslim/core/utils.py index 71a8168..33700db 100644 --- a/onnxslim/core/utils.py +++ b/onnxslim/core/utils.py @@ -20,6 +20,10 @@ def delete_node(node, input_var_idx=0, output_var_idx=0): if not isinstance(feed, (Variable, Constant)): feed.outputs.remove(node.inputs[input_var_idx]) feed.outputs.append(node.outputs[output_var_idx]) + for user in list(node.inputs[input_var_idx].outputs): + for i, input in enumerate(user.inputs): + if input == node.inputs[input_var_idx]: + user.inputs[i] = node.outputs[output_var_idx] node.outputs.clear() else: for next_node in next_nodes: