Felix Marty commited on
Commit
508744a
1 Parent(s): 3166d00
Files changed (1) hide show
  1. onnx_export.py +5 -2
onnx_export.py CHANGED
@@ -69,8 +69,11 @@ def convert_onnx(model_id: str, task: str, folder: str) -> List:
69
  if isinstance(atol, dict):
70
  atol = atol[task.replace("-with-past", "")]
71
 
72
- validate_model_outputs(onnx_config, model, output, onnx_outputs, atol)
73
- print(f"All good, model saved at: {output}")
 
 
 
74
 
75
  n_files = len([name for name in os.listdir(folder) if os.path.isfile(os.path.join(folder, name)) and not name.startswith(".")])
76
 
 
69
  if isinstance(atol, dict):
70
  atol = atol[task.replace("-with-past", "")]
71
 
72
+ try:
73
+ validate_model_outputs(onnx_config, model, output, onnx_outputs, atol)
74
+ print(f"All good, model saved at: {output}")
75
+ except ValueError:
76
+ print(f"An error occured, but the model was saved at: {args.output.as_posix()}")
77
 
78
  n_files = len([name for name in os.listdir(folder) if os.path.isfile(os.path.join(folder, name)) and not name.startswith(".")])
79