Merge pull request #364 from CBICA/sarthakpati-patch-1

Ensure ONNX-related print only happens once during training
mlcommons · Mar 22, 2022 · 7bb82d9 · 7bb82d9
2 parents 2ae23ff + e3578b9
commit 7bb82d9
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/GANDLF/utils/modelio.py b/GANDLF/utils/modelio.py
@@ -46,7 +46,9 @@ def save_model(model_dict, model, params, path, onnx_export=True):
     torch.save(model_dict, path)
 
     if onnx_export == False:
-        print("WARNING: Current model is not supported by ONNX/OpenVINO!")
+        if "onnx_print" not in params:
+            print("WARNING: Current model is not supported by ONNX/OpenVINO!")
+            params["onnx_print"] = True
         return
     else:
         try: