optimize argmax

5 years ago · 097db73d92
parent f3bc9b2f79
commit 097db73d92
1 changed files with 1 additions and 1 deletions
--- a/mindspore/lite/nnacl/arg_min_max.c
+++ b/mindspore/lite/nnacl/arg_min_max.c
@ -87,7 +87,7 @@ void ArgMinMaxTopknFp32(const float *input, float *output, const int *in_shape,
 }

 void ArgMinMax(const void *input, void *output, const int *in_shape, ArgMinMaxParameter *param) {
-  if (param->topk_ == 1 && !param->keep_dims_) {
+  if (param->topk_ == 1) {
    ArgMinMaxTopk1(input, output, in_shape, param);
    return;
  }