From a51bcf7bae608785fc6bb1ec9873200092fab8ab Mon Sep 17 00:00:00 2001 From: yefeng Date: Mon, 12 Jul 2021 14:43:51 +0800 Subject: [PATCH] fix infer --- .../backend/kernel_compiler/cpu/nnacl/infer/reduce_infer.c | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/reduce_infer.c b/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/reduce_infer.c index 30170c5c7d6..5058ab1c3ec 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/reduce_infer.c +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/reduce_infer.c @@ -70,7 +70,7 @@ int ReduceInferShape(const TensorC *const *inputs, size_t inputs_size, TensorC * bool keep_dims = param->keep_dims_; int out_shape[MAX_SHAPE_SIZE] = {0}; const size_t out_shape_size = 0; - if (inputs_size == 1) { + if (inputs_size == 1 || (inputs_size == 2 && inputs[1]->shape_size_ == 1 && inputs[1]->shape_[0] == 0)) { return ReduceOnAllAxes(input, output, out_shape, out_shape_size, keep_dims); } // get axes from input tensor