Fix large_model judgement (#1066)

intel · Jul 7, 2023 · 1b5337e · 1b5337e
1 parent c1ae63d
commit 1b5337e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/neural_compressor/model/onnx_model.py b/neural_compressor/model/onnx_model.py
@@ -44,7 +44,8 @@ def __init__(self, model, **kwargs):
         try:
             ort.InferenceSession(self._model.SerializeToString())
         except Exception as e:  # pragma: no cover
-            if 'maximum protobuf size of 2GB' in str(e) or 'string length exceeds max size' in str(e):
+            if 'maximum protobuf size of 2GB' in str(e) or 'string length exceeds max size' in str(e) or \
+                'protobuf parsing failed' in str(e):
                 self._is_large_model = True
                 if self._model_path is None:
                     logger.warning('Please use model path instead of onnx model object to quantize')