From c050880e5b5fd051a3430da42a77efa570835bd7 Mon Sep 17 00:00:00 2001 From: yuwenzho Date: Wed, 20 Sep 2023 15:05:58 +0800 Subject: [PATCH] fix typo Signed-off-by: yuwenzho --- neural_compressor/model/onnx_model.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/neural_compressor/model/onnx_model.py b/neural_compressor/model/onnx_model.py index a2169008c1c..ad031a29c96 100644 --- a/neural_compressor/model/onnx_model.py +++ b/neural_compressor/model/onnx_model.py @@ -65,7 +65,15 @@ def check_large_model(self): """Check model > 2GB.""" init_size = 0 for init in self._model.graph.initializer: - init_size += sys.getsizeof(init.SerializeToString()) + # if raise error of initializer size > 2GB, return True + try: + init_bytes = init.SerializeToString() + init_size += sys.getsizeof(init_bytes) + except Exception as e: + if "exceeds maximum protobuf size of 2GB" in str(e): + return True + else: + raise e if init_size > MAXIMUM_PROTOBUF: return True return False