fix bugs

PaddlePaddle · Sep 19, 2022 · a828143 · a828143
1 parent c2672f7
commit a828143
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 13 deletions.
diff --git a/cinn/frontend/net_builder_test.cc b/cinn/frontend/net_builder_test.cc
@@ -693,12 +693,12 @@ TEST(net_build, program_argmax_case1) {
         std::string line;
         for (int w = 0; w < W; ++w) {
           int index     = w + W * (h + H * (c + IN_C * n));
-          float in_data = input_data[index];
           int out_index = w + W * (h + H * n);
+          float in_data = input_data[index];
           int out_data  = output_data[out_index];
           EXPECT_LE(0, out_data);
           EXPECT_LT(out_data, IN_C);
-          int max_index   = w + W * (h + H * (out_data + OUT_C * n));
+          int max_index   = w + W * (h + H * (out_data + IN_C * n));
           float max_value = input_data[max_index];
           line += (std::to_string(out_data) + ", ");
           EXPECT_LE(in_data, max_value);
@@ -765,12 +765,12 @@ TEST(net_build, program_argmax_case2) {
         std::string line;
         for (int w = 0; w < W; ++w) {
           int index     = w + W * (h + H * (c + IN_C * n));
-          float in_data = input_data[index];
           int out_index = w + W * (h + H * n);
-          int out_data  = output_data[index];
+          float in_data = input_data[index];
+          int out_data  = output_data[out_index];
           EXPECT_LE(0, out_data);
           EXPECT_LT(out_data, IN_C);
-          int max_index   = w + W * (h + H * (out_data + n));
+          int max_index   = w + W * (h + H * (out_data + IN_C * n));
           float max_value = input_data[max_index];
           line += (std::to_string(out_data) + ", ");
           EXPECT_LE(in_data, max_value);
@@ -839,13 +839,12 @@ TEST(net_build, program_argmin_case1) {
         std::string line;
         for (int w = 0; w < W; ++w) {
           int index     = w + W * (h + H * (c + IN_C * n));
-          float in_data = input_data[index];
           int out_index = w + W * (h + H * n);
-          int out_data  = output_data[index];
-
+          float in_data = input_data[index];
+          int out_data  = output_data[out_index];
           EXPECT_LE(0, out_data);
           EXPECT_LT(out_data, IN_C);
-          int max_index   = w + W * (h + H * (out_data + OUT_C * n));
+          int max_index   = w + W * (h + H * (out_data + IN_C * n));
           float max_value = input_data[max_index];
           line += (std::to_string(out_data) + ", ");
           EXPECT_GE(in_data, max_value);
@@ -912,12 +911,12 @@ TEST(net_build, program_argmin_case2) {
         std::string line;
         for (int w = 0; w < W; ++w) {
           int index     = w + W * (h + H * (c + IN_C * n));
-          float in_data = input_data[index];
           int out_index = w + W * (h + H * n);
-          int out_data  = output_data[index];
+          float in_data = input_data[index];
+          int out_data  = output_data[out_index];
           EXPECT_LE(0, out_data);
           EXPECT_LT(out_data, IN_C);
-          int max_index   = w + W * (h + H * (out_data + n));
+          int max_index   = w + W * (h + H * (out_data + IN_C * n));
           float max_value = input_data[max_index];
           line += (std::to_string(out_data) + ", ");
           EXPECT_GE(in_data, max_value);

diff --git a/cinn/hlir/op/contrib/argmax.cc b/cinn/hlir/op/contrib/argmax.cc
@@ -70,7 +70,7 @@ Tensor Argmax(const Tensor &in_tensor,
     output_shape.push_back(Expr(1));
   }
 
-  auto sort_index = ArgSort(in_tensor, target, stages, pos_axis, true, name + "_index");
+  auto sort_index = ArgSort(in_tensor, target, stages, pos_axis, false, name + "_index");
   auto res        = Compute(
       output_shape,
       [=](const std::vector<Expr> &indices) {