From 405f04eb7c9fbca00e76f22c7310c677349f1edb Mon Sep 17 00:00:00 2001 From: Abdurrahman Akkas Date: Thu, 7 Aug 2025 16:10:40 -0700 Subject: [PATCH] Don't call erase if node is already erased in batch norm fusion. Summary: The erase call results in a warning log if the node is already erased by the dead code elimination earlier. Checking if the node is already erased so we don't unnecessarily pollute the log with warnings. Differential Revision: D79846881 --- torchao/quantization/pt2e/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchao/quantization/pt2e/utils.py b/torchao/quantization/pt2e/utils.py index 114f6b0ab4..41a26b62eb 100644 --- a/torchao/quantization/pt2e/utils.py +++ b/torchao/quantization/pt2e/utils.py @@ -758,7 +758,7 @@ def fold_bn_weights_into_conv_node( # since the node refers to a mutating op. Here we still need to call DCE first # to get rid of the unused getitem nodes that consume the BN node. m.graph.eliminate_dead_code() - if len(bn_node.users) == 0: + if not bn_node._erased and len(bn_node.users) == 0: m.graph.erase_node(bn_node)