Use nn_module_stack instead

Differential Revision: D85959415

Pull Request resolved: #3268

Author: tugsbayasgalan

torchao/quantization/pt2e/qat_utils.py

@@ -887,8 +887,8 @@ def _fold_conv_bn_qat(m: GraphModule) -> GraphModule:
             node.target == torch.ops.aten.add_.Tensor
             and node.args[0].op == "get_attr"
             and node.args[1] == 1
-            and torch.nn.modules.batchnorm.BatchNorm2d
-            in [val[1] for val in node.meta["source_fn_stack"]]
+            and "torch.nn.modules.batchnorm.BatchNorm2d"
+            in [val[1] for _, val in node.meta["nn_module_stack"].items()]
         ):
             m.graph.erase_node(node)