save exported program

_unittests/ut_torch_models/test_test_helpers.py

@@ -8,6 +8,7 @@
     ignore_warnings,
     requires_torch,
     requires_experimental,
+    requires_onnxscript,
 )
 from onnx_diagnostic.torch_models.test_helper import (
     get_inputs_for_task,
@@ -88,6 +89,7 @@ def test_validate_model_onnx_dynamo_ir(self):
         )
 
     @requires_torch("2.7")
+    @requires_onnxscript("0.4")
     @hide_stdout()
     @ignore_warnings(FutureWarning)
     def test_validate_model_onnx_dynamo_os_ort(self):

onnx_diagnostic/torch_models/test_helper.py

@@ -504,6 +504,7 @@ def validate_model(
                 print(f"[validate_model] -- dumps exported program in {dump_folder!r}...")
             with open(os.path.join(dump_folder, f"{folder_name}.ep"), "w") as f:
                 f.write(str(ep))
+            torch.export.save(ep, os.path.join(dump_folder, f"{folder_name}.pt2"))
             with open(os.path.join(dump_folder, f"{folder_name}.graph"), "w") as f:
                 f.write(str(ep.graph))
             if verbose:

onnx_diagnostic/torch_onnx/sbs.py

@@ -296,7 +296,7 @@ def post_process(obs):
             )
 
     for inp, v in zip(onx.graph.input, args):
-        onnx_results[inp.name] = v.numpy()
+        onnx_results[inp.name] = v.cpu().numpy()
         if verbose:
             print(
                 f"[run_aligned] +onnx-input: {inp.name}: "