Set save_as_external_data to True when saving a large model using onn…

…x.save (#3151) Signed-off-by: Huan Zhao <quic_huzh@quicinc.com>
quic · Jul 8, 2024 · b9f6f04 · b9f6f04
1 parent 76b0677
commit b9f6f04
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/TrainingExtensions/torch/src/python/aimet_torch/onnx_utils.py b/TrainingExtensions/torch/src/python/aimet_torch/onnx_utils.py
@@ -1580,7 +1580,8 @@ def save_initializer_restored_onnx_graph(original_model_path: str,
     """
     model = onnx.load(original_model_path)
     restored_model = restore_onnx_graph_initializers(model, inplace=True)
-    onnx.save(restored_model, restored_model_path)
+    save_as_external_data = model.ByteSize() >= onnx.checker.MAXIMUM_PROTOBUF
+    onnx.save(restored_model, restored_model_path, save_as_external_data=save_as_external_data)
 
 
 def restore_onnx_graph_initializers(model: onnx.ModelProto,