diff --git a/backends/arm/test/misc/test_extract_io_params_tosa.py b/backends/arm/test/misc/test_extract_io_params_tosa.py index 8483de63656..2afa3876081 100644 --- a/backends/arm/test/misc/test_extract_io_params_tosa.py +++ b/backends/arm/test/misc/test_extract_io_params_tosa.py @@ -60,16 +60,12 @@ def test_roundtrip_extracts_io_params(builder_method, quantizer_cls, partitioner operator_config = get_symmetric_quantization_config(is_qat=True) quantizer.set_global(operator_config) - exported = torch.export.export_for_training( - mod, copy.deepcopy(example_inputs), strict=True - ) + exported = torch.export.export(mod, copy.deepcopy(example_inputs), strict=True) prepared = prepare_pt2e(exported.module(), quantizer) _ = prepared(*example_inputs) converted = convert_pt2e(prepared) - final_export = torch.export.export_for_training( - converted, example_inputs, strict=True - ) + final_export = torch.export.export(converted, example_inputs, strict=True) partitioner = partitioner_cls(compile_spec) edge_prog = to_edge_transform_and_lower(final_export, partitioner=[partitioner]) diff --git a/backends/cortex_m/test/test_replace_quant_nodes.py b/backends/cortex_m/test/test_replace_quant_nodes.py index 08e75b17d9d..7d87bcb2b6a 100644 --- a/backends/cortex_m/test/test_replace_quant_nodes.py +++ b/backends/cortex_m/test/test_replace_quant_nodes.py @@ -1,5 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2025 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -16,7 +17,7 @@ ReplaceQuantNodesPass, ) from executorch.exir.dialects._ops import ops as exir_ops -from torch.export import export, export_for_training +from torch.export import export from torch.fx import GraphModule from torchao.quantization.pt2e.observer import HistogramObserver from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e @@ -125,9 +126,7 @@ def forward(self, x): example_inputs = (torch.randn(10, 11, 12),) # Step 1: Export and quantize the model - exported_model = export_for_training( - model.eval(), example_inputs, strict=True - ).module() + exported_model = export(model.eval(), example_inputs, strict=True).module() prepared_model = prepare_pt2e(exported_model, AddQuantizer()) quantized_model = convert_pt2e(prepared_model) diff --git a/docs/source/backends-arm-ethos-u.md b/docs/source/backends-arm-ethos-u.md index 71e3be105de..ce1dcdc6ea5 100644 --- a/docs/source/backends-arm-ethos-u.md +++ b/docs/source/backends-arm-ethos-u.md @@ -50,14 +50,14 @@ compile_spec = ArmCompileSpecBuilder().ethosu_compile_spec( ).build() # Post training quantization -graph_module = torch.export.export_for_training(mobilenet_v2, example_inputs).module() +graph_module = torch.export.export(mobilenet_v2, example_inputs).module() quantizer = EthosUQuantizer(compile_spec) operator_config = get_symmetric_quantization_config(is_per_channel=False) quantizer.set_global(operator_config) graph_module = prepare_pt2e(graph_module, quantizer) graph_module(*example_inputs) graph_module = convert_pt2e(graph_module) -exported_program = torch.export.export_for_training(graph_module, example_inputs) +exported_program = torch.export.export(graph_module, example_inputs) # Lower the exported program to the Ethos-U backend and save pte file. edge_program_manager = to_edge_transform_and_lower( diff --git a/examples/arm/aot_arm_compiler.py b/examples/arm/aot_arm_compiler.py index d6a1eab3205..92573dad55c 100644 --- a/examples/arm/aot_arm_compiler.py +++ b/examples/arm/aot_arm_compiler.py @@ -702,7 +702,7 @@ def quantize_model(args, model: torch.nn.Module, example_inputs, compile_spec): args.evaluate_config, ) # Wrap quantized model back into an exported_program - exported_program = torch.export.export_for_training( + exported_program = torch.export.export( model_int8, example_inputs, strict=args.strict_export ) @@ -795,9 +795,9 @@ def transform_for_cortex_m_backend(edge): ) model = original_model.eval() - # export_for_training under the assumption we quantize, the exported form also works + # export under the assumption we quantize, the exported form also works # in to_edge if we don't quantize - exported_program = torch.export.export_for_training( + exported_program = torch.export.export( model, example_inputs, strict=args.strict_export ) model = exported_program.module() diff --git a/examples/arm/ethos_u_minimal_example.ipynb b/examples/arm/ethos_u_minimal_example.ipynb index 72caed50149..d1f27a95d64 100644 --- a/examples/arm/ethos_u_minimal_example.ipynb +++ b/examples/arm/ethos_u_minimal_example.ipynb @@ -58,7 +58,7 @@ "\n", "model = Add()\n", "model = model.eval()\n", - "exported_program = torch.export.export_for_training(model, example_inputs)\n", + "exported_program = torch.export.export(model, example_inputs)\n", "graph_module = exported_program.module()\n", "\n", "_ = graph_module.print_readable()" @@ -112,7 +112,7 @@ "_ = quantized_graph_module.print_readable()\n", "\n", "# Create a new exported program using the quantized_graph_module\n", - "quantized_exported_program = torch.export.export_for_training(quantized_graph_module, example_inputs)" + "quantized_exported_program = torch.export.export(quantized_graph_module, example_inputs)" ] }, {