From d5fecc8535aaf92e99856336870ff9a90e75fab2 Mon Sep 17 00:00:00 2001 From: zhaozhenlong Date: Sat, 19 Sep 2020 11:05:02 +0800 Subject: [PATCH] fix malloc failed handle --- mindspore/lite/src/ops/batch_norm.cc | 9 +++++---- mindspore/lite/src/populate_parameter.cc | 5 +++++ 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/mindspore/lite/src/ops/batch_norm.cc b/mindspore/lite/src/ops/batch_norm.cc index 736e6a9441f..5ae19f1a761 100644 --- a/mindspore/lite/src/ops/batch_norm.cc +++ b/mindspore/lite/src/ops/batch_norm.cc @@ -38,12 +38,13 @@ int BatchNorm::UnPackAttr(const Primitive &prim, const std::vector & } if (this->primitive_->value.value == nullptr) { auto attr = new (std::nothrow) schema::FusedBatchNormT(); - attr->epsilon = GetValue(prim.GetAttr("epsilon")); - this->primitive_->value.value = attr; - if (this->primitive_->value.value == nullptr) { - MS_LOG(ERROR) << "new primitiveT value failed"; + if (attr == nullptr) { + MS_LOG(ERROR) << "new FusedBatchNormT failed"; + delete this->primitive_; return RET_ERROR; } + attr->epsilon = GetValue(prim.GetAttr("epsilon")); + this->primitive_->value.value = attr; } return RET_OK; } diff --git a/mindspore/lite/src/populate_parameter.cc b/mindspore/lite/src/populate_parameter.cc index b7ff742dc8f..c52b403dfae 100644 --- a/mindspore/lite/src/populate_parameter.cc +++ b/mindspore/lite/src/populate_parameter.cc @@ -1536,6 +1536,11 @@ OpParameter *PopulateL2NormParameter(const mindspore::lite::PrimitiveC *primitiv auto axis_vec = param->GetAxis(); l2_norm_parameter->axis_num_ = axis_vec.size(); l2_norm_parameter->axis_ = reinterpret_cast(malloc(axis_vec.size() * sizeof(int))); + if (l2_norm_parameter->axis_ == nullptr) { + MS_LOG(ERROR) << "malloc axis_ data failed"; + free(l2_norm_parameter); + return nullptr; + } for (size_t i = 0; i < axis_vec.size(); i++) { l2_norm_parameter->axis_[i] = axis_vec[i]; }