Switch XNNPack tests to use export_for_training

Differential Revision: D61684468

Pull Request resolved: #4867

backends/xnnpack/test/test_xnnpack_utils.py

@@ -72,6 +72,7 @@
     get_symmetric_quantization_config,
     XNNPACKQuantizer,
 )
+from torch.export import export_for_training
 
 from torch.testing import FileCheck
 
@@ -315,10 +316,11 @@ def quantize_and_test_model_with_quantizer(
     ):
         module.eval()
         # program capture
-        m = torch._export.capture_pre_autograd_graph(
+
+        m = export_for_training(
             module,
             example_inputs,
-        )
+        ).module()
 
         quantizer = XNNPACKQuantizer()
         quantization_config = get_symmetric_quantization_config()

backends/xnnpack/test/tester/tester.py

@@ -14,7 +14,6 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
 
 import torch
-import torch.export._trace as export_trace
 from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
 from executorch.backends.xnnpack.passes import XNNPACKPassManager
 from executorch.backends.xnnpack.utils.configs import get_xnnpack_edge_compile_config
@@ -31,6 +30,7 @@
 from executorch.exir.passes.sym_shape_eval_pass import ConstraintBasedSymShapeEvalPass
 
 from executorch.exir.print_program import pretty_print, print_program
+from torch.export import export_for_training
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
@@ -157,10 +157,10 @@ def __init__(
     def run(
         self, artifact: torch.nn.Module, inputs: Optional[Tuple[torch.Tensor]]
     ) -> None:
-        captured_graph = export_trace._export(
-            artifact, inputs, pre_dispatch=True
-        ).module()
+        assert inputs is not None
+        captured_graph = export_for_training(artifact, inputs).module()
 
+        assert isinstance(captured_graph, torch.fx.GraphModule)
         prepared = prepare_pt2e(captured_graph, self.quantizer)
 
         if self.calibrate: