[torch][ao] Use returned model from Quantizer.transform_for_annotation in prepare_pt2e (pytorch#132893)

Summary:
The Quantizer subclass can return a new model from `transform_for_annotation`,
and this is common if it uses any ExportPass subclass which does not mutate in-place.

Use the returned model instead of assuming its the same.

Differential Revision: D60869676

Pull Request resolved: pytorch#132893
Approved by: https://github.com/jerryzh168

Author: dulinriley

test/quantization/pt2e/test_quantize_pt2e.py

@@ -1500,9 +1500,14 @@ class TestQuantizer(Quantizer):
             def transform_for_annotation(
                 self, model: torch.fx.GraphModule
             ) -> torch.fx.GraphModule:
-                for n in model.graph.nodes:
+                # Make a copy of the graph to ensure that we are using the
+                # return value of this function.
+                graph = torch.fx.Graph()
+                graph.graph_copy(model.graph, {})
+                for n in graph.nodes:
                     if n.target == torch.ops.aten.add.Tensor:
                         n.target = torch.ops.aten.mul.Tensor
+                model = torch.fx.GraphModule(model, graph)
                 return model
 
             def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule:

torch/ao/quantization/quantize_pt2e.py

@@ -96,7 +96,7 @@ def calibrate(model, data_loader):
     # to be quantized before fusion
     # TODO: (maybe) rewrite this with subgraph_rewriter
     _fuse_conv_bn_(model)
-    quantizer.transform_for_annotation(model)
+    model = quantizer.transform_for_annotation(model)
     quantizer.annotate(model)
     quantizer.validate(model)
     model = prepare(model, node_name_to_scope, is_qat=False)
@@ -165,7 +165,7 @@ def train_loop(model, train_data):
     torch._C._log_api_usage_once("quantization_api.quantize_pt2e.prepare_qat_pt2e")
     original_graph_meta = model.meta
     node_name_to_scope = _get_node_name_to_scope(model)
-    quantizer.transform_for_annotation(model)
+    model = quantizer.transform_for_annotation(model)
     quantizer.annotate(model)
     quantizer.validate(model)
     # Perform fusion after annotate to avoid quantizing ops in the new