Browse Source

!8371 fix infoshape problems of pooling ops

From: @liuwenhao4
Reviewed-by: @zhang_xue_tong,@zhanghaibo5
Signed-off-by: @zhang_xue_tong
tags/v1.1.0
mindspore-ci-bot Gitee 5 years ago
parent
commit
0c46dbc024
1 changed files with 2 additions and 2 deletions
  1. +2
    -2
      mindspore/lite/src/ops/pooling.cc

+ 2
- 2
mindspore/lite/src/ops/pooling.cc View File

@@ -228,8 +228,8 @@ int Pooling::InferShape(std::vector<Tensor *> inputs_, std::vector<Tensor *> 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;
}


Loading…
Cancel
Save