diff --git a/mindspore/lite/src/ops/pooling.cc b/mindspore/lite/src/ops/pooling.cc index 3b415339c6..b8cfc4fd46 100644 --- a/mindspore/lite/src/ops/pooling.cc +++ b/mindspore/lite/src/ops/pooling.cc @@ -229,8 +229,8 @@ int Pooling::InferShape(std::vector inputs_, std::vector out } } auto input_shape = input->shape(); - input_shape.at(1) = output_h; - input_shape.at(2) = output_w; + input_shape.at(1) = output_h > 0 ? output_h : 1; + input_shape.at(2) = output_w > 0 ? output_w : 1; output->set_shape(input_shape); return RET_OK; }