This commit is contained in:
CHEN YU 2022-11-02 17:15:40 +08:00
parent ee4b33ad62
commit e520a7687c
4 changed files with 17 additions and 2 deletions

View File

@ -23,9 +23,9 @@ int GatherDInferShape(const TensorC *const *inputs, size_t inputs_size, TensorC
if (ret != NNACL_OK) {
return ret;
}
const int input_size_limit = 2;
const int input_size_limit = 3;
const int output_size_limit = 1;
if (inputs_size < input_size_limit || outputs_size != output_size_limit) {
if (inputs_size != input_size_limit || outputs_size != output_size_limit) {
return NNACL_ERR;
}
const TensorC *input = inputs[0];

View File

@ -19,6 +19,11 @@
int LogSoftmaxInferShape(const TensorC *const *inputs, size_t inputs_size, TensorC **outputs, size_t outputs_size,
OpParameter *parameter) {
const int input_size_limit = 1;
const int output_size_limit = 1;
if (inputs_size != input_size_limit || outputs_size != output_size_limit) {
return NNACL_ERR;
}
int check_ret = CheckAugmentWithMinSize(inputs, inputs_size, outputs, outputs_size, parameter, 1, 1);
if (check_ret != NNACL_OK) {
return check_ret;

View File

@ -51,6 +51,10 @@ void Init_MSC_VER_kernels(void) {
bool SupportKernelC(int opType, int format, int dataType) {
Init_MSC_VER_kernels();
const int length = 16;
if (REGIST_DT(dataType) < 0 || REGIST_DT(dataType) >= length) {
return false;
}
KernelCreator creator = g_kernelCreatorRegistry[opType][format][REGIST_DT(dataType)];
return creator != NULL;
}

View File

@ -35,6 +35,12 @@ OpParameter *PopulateArgMaxParameter(const void *prim) {
free(arg_param);
return nullptr;
}
if (param->axis() > INT32_MAX) {
MS_LOG(ERROR) << "check le fail, value1:" << param->axis() << " value2: " << INT32_MAX;
free(arg_param);
return nullptr;
}
arg_param->axis_ = param->axis();
arg_param->topk_ = param->top_k();
arg_param->out_value_ = param->out_max_value();