aot_arm_compiler: Add argument to opt out of strict exports (pytorch#12750)

This way models that fail strict exporting may still be exported, while
still maintaining the preferred, strict flow.


cc @digantdesai @freddan80 @per @zingo @oscarandersson8218

Co-authored-by: Tobias Bladh <[email protected]>

Author: 2 people

examples/arm/aot_arm_compiler.py

@@ -580,6 +580,13 @@ def get_args():
         default="Arm/vela.ini",
         help="Specify custom vela configuration file (vela.ini)",
     )
+    parser.add_argument(
+        "--non_strict_export",
+        dest="strict_export",
+        required=False,
+        action="store_false",
+        help="Disable strict checking while exporting models.",
+    )
     args = parser.parse_args()
 
     if args.evaluate and (
@@ -696,7 +703,7 @@ def quantize_model(args, model: torch.nn.Module, example_inputs, compile_spec):
     )
     # Wrap quantized model back into an exported_program
     exported_program = torch.export.export_for_training(
-        model_int8, example_inputs, strict=True
+        model_int8, example_inputs, strict=args.strict_export
     )
 
     return model_int8, exported_program
@@ -791,7 +798,7 @@ def transform_for_cortex_m_backend(edge):
     # export_for_training 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(
-        model, example_inputs, strict=True
+        model, example_inputs, strict=args.strict_export
     )
     model = exported_program.module()
     model_fp32 = model