Skip to content

Commit 7cb920b

Browse files
authored
Don't call erase if node is already erased in batch norm fusion.
Differential Revision: D79846881 Pull Request resolved: #2716
1 parent 0315628 commit 7cb920b

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

torchao/quantization/pt2e/utils.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -758,7 +758,7 @@ def fold_bn_weights_into_conv_node(
758758
# since the node refers to a mutating op. Here we still need to call DCE first
759759
# to get rid of the unused getitem nodes that consume the BN node.
760760
m.graph.eliminate_dead_code()
761-
if len(bn_node.users) == 0:
761+
if not bn_node._erased and len(bn_node.users) == 0:
762762
m.graph.erase_node(bn_node)
763763

764764

0 commit comments

Comments
 (0)