save ep

Author: xadupre

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}: "