diff --git a/onnxruntime/core/graph/model.cc b/onnxruntime/core/graph/model.cc index f03e311168..ca31a76525 100644 --- a/onnxruntime/core/graph/model.cc +++ b/onnxruntime/core/graph/model.cc @@ -237,7 +237,7 @@ Status Model::Load(std::istream& model_istream, ModelProto* p_model_proto) { if (!p_model_proto) { return Status(ONNXRUNTIME, INVALID_ARGUMENT, "Null model_proto ptr."); } - google::protobuf::io::IstreamInputStream zero_copy_input(&model_istream); + google::protobuf::io::IstreamInputStream zero_copy_input(&model_istream, 1 << 20); const bool result = p_model_proto->ParseFromZeroCopyStream(&zero_copy_input) && model_istream.eof(); if (!result) { return Status(ONNXRUNTIME, INVALID_PROTOBUF, "Failed to load model because protobuf parsing failed."); @@ -447,7 +447,8 @@ Status Model::Load(int fd, ONNX_NAMESPACE::ModelProto& model_proto) { } #if GOOGLE_PROTOBUF_VERSION >= 3002000 - const bool result = model_proto.ParseFromFileDescriptor(fd); + FileInputStream input(fd, 1 << 20); + const bool result = model_proto.ParseFromZeroCopyStream(&input) && input.GetErrno() == 0; if (!result) { return Status(ONNXRUNTIME, INVALID_PROTOBUF, "Protobuf parsing failed."); } diff --git a/onnxruntime/test/onnx/TestCase.cc b/onnxruntime/test/onnx/TestCase.cc index 933482dd2f..7d6ed4a00a 100644 --- a/onnxruntime/test/onnx/TestCase.cc +++ b/onnxruntime/test/onnx/TestCase.cc @@ -176,7 +176,9 @@ class OnnxModelInfo : public TestModelInfo { } ONNX_NAMESPACE::ModelProto model_pb; - if (!model_pb.ParseFromFileDescriptor(model_fd)) { + ::google::protobuf::io::FileInputStream input(model_fd, 1 << 20); + const bool parse_result = model_pb.ParseFromZeroCopyStream(&input) && input.GetErrno() == 0; + if (!parse_result) { (void)Env::Default().FileClose(model_fd); ORT_THROW("Failed to load model because protobuf parsing failed."); }