From 1b5337ef0a439eeab4d6ca1ada8f7ed4c8e14b9c Mon Sep 17 00:00:00 2001 From: "Wang, Mengni" Date: Fri, 7 Jul 2023 13:35:44 +0800 Subject: [PATCH] Fix large_model judgement (#1066) --- neural_compressor/model/onnx_model.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/neural_compressor/model/onnx_model.py b/neural_compressor/model/onnx_model.py index 27583181a0a..86027aece40 100644 --- 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')