diff --git a/examples/xnnpack/aot_compiler.py b/examples/xnnpack/aot_compiler.py index c3538db4d8..e1542245ac 100644 --- a/examples/xnnpack/aot_compiler.py +++ b/examples/xnnpack/aot_compiler.py @@ -92,6 +92,7 @@ logging.info("Quantizing Model...") # TODO(T165162973): This pass shall eventually be folded into quantizer model = quantize(model, example_inputs) + ep = torch.export.export_for_training(model, example_inputs) edge = to_edge_transform_and_lower( ep,