Browse Source

fix infoshape problems of pooling ops

tags/v1.1.0
liuwenhao4 5 years ago
parent
commit
2a8dc5dd27
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

@@ -229,8 +229,8 @@ int Pooling::InferShape(std::vector<Tensor *> inputs_, std::vector<Tensor *> out
} }
} }
auto input_shape = input->shape(); 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); output->set_shape(input_shape);
return RET_OK; return RET_OK;
} }


Loading…
Cancel
Save