diff --git a/mindspore/lite/src/ops/populate/space_to_batch_nd_populate.cc b/mindspore/lite/src/ops/populate/space_to_batch_nd_populate.cc index 426718c1ed..ce63189155 100644 --- a/mindspore/lite/src/ops/populate/space_to_batch_nd_populate.cc +++ b/mindspore/lite/src/ops/populate/space_to_batch_nd_populate.cc @@ -34,12 +34,14 @@ OpParameter *PopulateSpaceToBatchNDParameter(const mindspore::lite::PrimitiveC * space_batch_param_nd->m_ = block_sizes.size(); if (block_sizes.size() > std::numeric_limits::max() / sizeof(int)) { MS_LOG(ERROR) << "The value of block_sizes.size() is too big"; + free(space_batch_param_nd); return nullptr; } memcpy(space_batch_param_nd->block_sizes_, (block_sizes.data()), block_sizes.size() * sizeof(int)); auto paddings = ((mindspore::lite::SpaceToBatchND *)primitive)->GetPaddings(); if (paddings.size() > std::numeric_limits::max() / sizeof(int)) { MS_LOG(ERROR) << "The value of paddings.size() is too big"; + free(space_batch_param_nd); return nullptr; } memcpy(space_batch_param_nd->paddings_, (paddings.data()), paddings.size() * sizeof(int)); diff --git a/mindspore/lite/src/ops/populate/space_to_batch_populate.cc b/mindspore/lite/src/ops/populate/space_to_batch_populate.cc index 12a91ee7ee..d4d803f3b9 100644 --- a/mindspore/lite/src/ops/populate/space_to_batch_populate.cc +++ b/mindspore/lite/src/ops/populate/space_to_batch_populate.cc @@ -36,12 +36,14 @@ OpParameter *PopulateSpaceToBatchParameter(const mindspore::lite::PrimitiveC *pr space_batch_param->m_ = block_sizes.size(); if (block_sizes.size() > std::numeric_limits::max() / sizeof(int)) { MS_LOG(ERROR) << "The value of block_sizes.size() is too big"; + free(space_batch_param); return nullptr; } memcpy(space_batch_param->block_sizes_, (block_sizes.data()), block_sizes.size() * sizeof(int)); auto paddings = ((mindspore::lite::SpaceToBatch *)primitive)->Paddings(); if (paddings.size() > std::numeric_limits::max() / sizeof(int)) { MS_LOG(ERROR) << "The value of paddings.size() is too big"; + free(space_batch_param); return nullptr; } memcpy(space_batch_param->paddings_, (paddings.data()), paddings.size() * sizeof(int)); diff --git a/mindspore/lite/src/ops/populate/strided_slice_populate.cc b/mindspore/lite/src/ops/populate/strided_slice_populate.cc index efe6dfe385..441843518f 100644 --- a/mindspore/lite/src/ops/populate/strided_slice_populate.cc +++ b/mindspore/lite/src/ops/populate/strided_slice_populate.cc @@ -37,24 +37,28 @@ OpParameter *PopulateStridedSliceParameter(const mindspore::lite::PrimitiveC *pr auto begin = ((lite::StridedSlice *)primitive)->GetBegins(); if (begin.size() > std::numeric_limits::max() / sizeof(int)) { MS_LOG(ERROR) << "The value of begin.size() is too big"; + free(strided_slice_param); return nullptr; } memcpy(strided_slice_param->begins_, (begin.data()), begin.size() * sizeof(int)); auto end = ((lite::StridedSlice *)primitive)->GetEnds(); if (end.size() > std::numeric_limits::max() / sizeof(int)) { MS_LOG(ERROR) << "The value of end.size() is too big"; + free(strided_slice_param); return nullptr; } memcpy(strided_slice_param->ends_, (end.data()), end.size() * sizeof(int)); auto stride = ((lite::StridedSlice *)primitive)->GetStrides(); if (stride.size() > std::numeric_limits::max() / sizeof(int)) { MS_LOG(ERROR) << "The value of stride.size() is too big"; + free(strided_slice_param); return nullptr; } memcpy(strided_slice_param->strides_, (stride.data()), stride.size() * sizeof(int)); auto in_shape = ((lite::StridedSlice *)primitive)->GetInShape(); if (in_shape.size() > std::numeric_limits::max() / sizeof(int)) { MS_LOG(ERROR) << "The value of in_shape.size() is too big"; + free(strided_slice_param); return nullptr; } memcpy(strided_slice_param->in_shape_, (in_shape.data()), in_shape.size() * sizeof(int)); diff --git a/mindspore/lite/src/ops/space_to_depth.cc b/mindspore/lite/src/ops/space_to_depth.cc index a261d006fe..764a332308 100644 --- a/mindspore/lite/src/ops/space_to_depth.cc +++ b/mindspore/lite/src/ops/space_to_depth.cc @@ -63,13 +63,13 @@ int SpaceToDepth::InferShape(std::vector inputs, std::vectorprimitive_ != nullptr); if (outputs.size() != kSpaceToDepthOutputNum || inputs.size() != kSpaceToDepthInputNum) { MS_LOG(ERROR) << "Invalid output/input size! output size: " << outputs.size() << ",input size: " << inputs.size(); - return 1; + return RET_ERROR; } auto input = inputs.at(0); if (input->format() != schema::Format::Format_NHWC) { MS_LOG(ERROR) << "space_to_depth only support NHWC now!"; - return 1; + return RET_ERROR; } outputs.at(0)->set_format(input->format()); outputs.at(0)->set_data_type(input->data_type()); @@ -79,14 +79,18 @@ int SpaceToDepth::InferShape(std::vector inputs, std::vectorshape(); if (input_shape.size() != kDimension_4d) { MS_LOG(ERROR) << "input shape dimension size should == " << kDimension_4d; - return 1; + return RET_ERROR; } int32_t block_size = GetBlockSize(); + if (block_size == 0) { + MS_LOG(ERROR) << "block_size is zero"; + return RET_ERROR; + } if (input_shape.at(NHWC_H) % block_size != 0 || input_shape.at(NHWC_H) == 0 || input_shape.at(NHWC_W) % block_size != 0 || input_shape.at(NHWC_W) == 0) { MS_LOG(ERROR) << "input dimension h or w size error!"; - return 1; + return RET_ERROR; } std::vector output_shape(input_shape.size()); output_shape.at(NHWC_N) = input_shape.at(NHWC_N);