Skip onnx.check_model if model is too large

neuralmagic · Apr 10, 2023 · f0ba51c · f0ba51c
1 parent 9382d2f
commit f0ba51c
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/src/sparseml/onnx/utils/helpers.py b/src/sparseml/onnx/utils/helpers.py
@@ -87,7 +87,14 @@ def validate_onnx_file(path: str):
     """
     try:
         onnx_model = check_load_model(path)
-        onnx.checker.check_model(onnx_model)
+
+        if onnx_model.ByteSize() < onnx.checker.MAXIMUM_PROTOBUF:
+            onnx.checker.check_model(onnx_model)
+        else:
+            _LOGGER.warning(
+                "onnx check_model skipped as model exceeds maximum protobuf size of 2GB"
+            )
+
         if not onnx_model.opset_import:
             raise ValueError("could not parse opset_import")
     except Exception as err: