未验证 提交 8e85729c 编写于 作者: H HappyAngel 提交者: GitHub

change axis_size = 4 kernel choose, test=develop

上级 ee1c6d91
......@@ -34,7 +34,7 @@ void SoftmaxCompute::Run() {
int inner_num = x_dims.Slice(axis + 1, x_rank).production();
int axis_size = x_dims[axis];
if (inner_num == 1) {
if (axis_size >= 4) {
if (axis_size > 4) {
lite::arm::math::softmax_inner1_large_axis(
din, dout, outer_num, axis_size);
} else {
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册