update

bytedance · Jul 27, 2024 · d803418 · d803418
1 parent 3e39320
commit d803418
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/numerical_test/mlir_tests/cpu_ops/custom_call_byteir_arg_max.mlir b/tests/numerical_test/mlir_tests/cpu_ops/custom_call_byteir_arg_max.mlir
@@ -1,4 +1,4 @@
-func.func @byteir.arg_max$return_2(%arg0: tensor<3x4xf32>) -> (tensor<3xf32>, tensor<3xi64>) {
-  %0:2 = stablehlo.custom_call @byteir.arg_max(%arg0) {byteir_attrs = {axis = 1 : i64, keep_dims = false, select_last_index = false}} : (tensor<3x4xf32>) -> (tensor<3xf32>, tensor<3xi64>)
+func.func @byteir.arg_max$return_2(%arg0: tensor<3x128xf32>) -> (tensor<3xf32>, tensor<3xi64>) {
+  %0:2 = stablehlo.custom_call @byteir.arg_max(%arg0) {byteir_attrs = {axis = 1 : i64, keep_dims = false, select_last_index = false}} : (tensor<3x128xf32>) -> (tensor<3xf32>, tensor<3xi64>)
   return %0#0, %0#1 : tensor<3xf32>, tensor<3xi64>
 }