Remove torch.export.export_for_inference

Differential Revision: D71069057

Pull Request resolved: #1877

Author: gmagogsfm

examples/sam2_amg_server/compile_export_utils.py

@@ -119,9 +119,10 @@ def aot_compile(
             "triton.cudagraphs": True,
         }
 
-    from torch.export import export_for_inference
+    from torch.export import export_for_training
 
-    exported = export_for_inference(fn, sample_args, sample_kwargs)
+    exported = export_for_training(fn, sample_args, sample_kwargs, strict=True)
+    exported.run_decompositions()
     output_path = torch._inductor.aoti_compile_and_package(
         exported,
         package_path=str(path),

examples/sam2_vos_example/compile_export_utils.py

@@ -82,9 +82,10 @@ def aot_compile(
             "triton.cudagraphs": True,
         }
 
-    from torch.export import export_for_inference
+    from torch.export import export_for_training
 
-    exported = export_for_inference(fn, sample_args, sample_kwargs)
+    exported = export_for_training(fn, sample_args, sample_kwargs, strict=True)
+    exported.run_decompositions()
     output_path = torch._inductor.aoti_compile_and_package(
         exported,
         package_path=str(path),