Remove export_for_training

Differential Revision: D90114683

Pull Request resolved: #16434

Author: tugsbayasgalan

.ci/scripts/wheel/test_windows.py

@@ -30,9 +30,7 @@ def test_model_xnnpack(model: Model, quantize: bool) -> None:
 
     if quantize:
         quant_type = MODEL_NAME_TO_OPTIONS[str(model)].quantization
-        model_instance = torch.export.export_for_training(
-            model_instance, example_inputs
-        )
+        model_instance = torch.export.export(model_instance, example_inputs)
         model_instance = quantize_xnn(
             model_instance.module(), example_inputs, quant_type
         )

backends/samsung/utils/export_utils.py

@@ -63,7 +63,7 @@ def quantize_module(
     quantizer = EnnQuantizer()
     quantizer.setup_quant_params(precision, is_per_channel, is_qat)
     logging.info("Export nn module for quantization...")
-    exported_module = torch.export.export_for_training(module, inputs).module()
+    exported_module = torch.export.export(module, inputs).module()
     DecomposeScaledDotProductAttention()(exported_module)
     logging.info("Quantizing the module...")
     annotated_module = prepare_pt2e(exported_module, quantizer)

examples/mediatek/model_export_scripts/gemma.py

@@ -393,7 +393,7 @@ def export_to_et_ir(
         max_num_token, max_cache_size, True
     )
     print("Getting pre autograd ATen Dialect Graph")
-    pre_autograd_aten_dialect = torch.export.export_for_training(
+    pre_autograd_aten_dialect = torch.export.export(
         model, example_inputs, dynamic_shapes=dynamic_shapes, strict=True
     ).module()  # NOTE: Will be replaced with export
     quantizer = NeuropilotQuantizer()

examples/mediatek/model_export_scripts/phi.py

@@ -338,7 +338,7 @@ def export_to_et_ir(
         max_num_token, max_cache_size, True
     )
     print("Getting pre autograd ATen Dialect Graph")
-    pre_autograd_aten_dialect = torch.export.export_for_training(
+    pre_autograd_aten_dialect = torch.export.export(
         model, example_inputs, dynamic_shapes=dynamic_shapes, strict=True
     ).module()  # NOTE: Will be replaced with export
     quantizer = NeuropilotQuantizer()

examples/mediatek/model_export_scripts/qwen.py

@@ -339,7 +339,7 @@ def export_to_et_ir(
         max_num_token, max_cache_size, True
     )
     print("Getting pre autograd ATen Dialect Graph")
-    pre_autograd_aten_dialect = torch.export.export_for_training(
+    pre_autograd_aten_dialect = torch.export.export(
         model, example_inputs, dynamic_shapes=dynamic_shapes, strict=True
     ).module()  # NOTE: Will be replaced with export
     quantizer = NeuropilotQuantizer()

examples/mediatek/model_export_scripts/whisper.py

@@ -410,7 +410,7 @@ def export_to_et_ir(
         max_num_token, max_cache_size, True
     )
     print("Getting pre autograd ATen Dialect Graph")
-    pre_autograd_aten_dialect = torch.export.export_for_training(
+    pre_autograd_aten_dialect = torch.export.export(
         model, example_inputs, dynamic_shapes=dynamic_shapes, strict=True
     ).module()  # NOTE: Will be replaced with export
     quantizer = NeuropilotQuantizer()
@@ -483,7 +483,7 @@ def export_encoder_to_et_ir(
     print(f"Exporting Encoder to PTE")
     example_inputs = model.get_example_inputs(num_mel_bins)
     print("Getting pre autograd ATen Dialect Graph")
-    pre_autograd_aten_dialect = torch.export.export_for_training(
+    pre_autograd_aten_dialect = torch.export.export(
         model, example_inputs, strict=True
     ).module()  # NOTE: Will be replaced with export
     quantizer = NeuropilotQuantizer()

exir/tests/test_dim_order_utils.py

@@ -44,6 +44,6 @@ def forward(self, t1, t2):
         y = torch.randn(5, 6)
         M(x, y)
 
-        expo_prog = torch.export.export_for_training(M, (x, y))
+        expo_prog = torch.export.export(M, (x, y))
         edge_prog = to_edge_transform_and_lower(expo_prog)
         edge_prog.to_executorch()