fix typo

Signed-off-by: yuwenzho <[email protected]>
intel · Sep 21, 2023 · c050880 · c050880
1 parent 8a9a8aa
commit c050880
Showing 1 changed file with 9 additions and 1 deletion.
diff --git 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