提交 052c5de5 编写于 作者: Y yejianwu

check keep_dims for ReduceMeanOp in GPU

上级 fbc1d019
......@@ -56,7 +56,7 @@ MaceStatus ReduceMeanKernel<T>::Compute(
const Tensor *input,
Tensor *output) {
MACE_CHECK_NOTNULL(input);
// MACE_CHECK(keep_dims_, "reduce mean gpu only support keep dims.");
MACE_CHECK(keep_dims_, "reduce mean gpu only support keep dims.");
MACE_CHECK(input->dim_size() == 4,
"reduce mean gpu only support 4-dim input");
MACE_CHECK(axis_.size() == 2 && axis_[0] == 1 && axis_[1] == 2,
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册