Fix Argmax Layout autotune (#44080)

paddle/fluid/imperative/layout_transformer.h

@@ -374,8 +374,9 @@ class ArgmaxOpTransformer
     bool keep_dims = BOOST_GET_CONST(bool, (*attrs)["keepdims"]);
     if (keep_dims) {
       if (var_layout != DataLayout::UNDEFINED) {
-        std::vector<int> perm_nhwc = {0, 2, 3, 1};
-        std::vector<int> perm_nchw = {0, 3, 1, 2};
+        std::vector<int> perm_nhwc = {0, 3, 1, 2};
+        std::vector<int> perm_nchw = {0, 2, 3, 1};
+
         auto perm = var_layout == DataLayout::NHWC ? perm_nhwc : perm_nchw;
         switch (AttrTypeID((*attrs)["axis"])) {
           case paddle::framework::proto::AttrType::INT: {

python/paddle/fluid/tests/unittests/test_layout_autotune.py

@@ -146,7 +146,7 @@ class LayoutAutoTune(unittest.TestCase):
             out = paddle.argmax(conv_out, axis=1, keepdim=True)
 
         self.assertEqual(conv_out.shape, [1, 14, 12, 8])
-        self.assertEqual(out.shape, [1, 14, 1, 8])
+        self.assertEqual(out.shape, [1, 14, 12, 1])
 
     def test_argmax_op_transposer(self):
         if not self.use_autoune():