diff --git a/mace/ops/BUILD b/mace/ops/BUILD
index 7d5e68322e52753e54858a9c6f2dd985bb96e466..89f6b8b78bb1f31a912817140756d562275a61ce 100644
--- a/mace/ops/BUILD
+++ b/mace/ops/BUILD
@@ -61,22 +61,6 @@ cc_test(
     ],
 )
 
-cc_test(
-    name = "concat_test",
-    testonly = 1,
-    srcs = glob(
-        ["concat_test.cc"],
-    ),
-    copts = ["-std=c++11"],
-    linkopts = ["-fopenmp"],
-    linkstatic = 1,
-    deps = [
-        ":ops",
-        ":test",
-        "@gtest//:gtest_main",
-    ],
-)
-
 cc_test(
     name = "ops_benchmark",
     testonly = 1,
diff --git a/mace/ops/concat_benchmark.cc b/mace/ops/concat_benchmark.cc
index 01db5a43bcc3a493d63d4fafd90be6e1d868d295..0e4425251426e8056a2973352ded79cd2b4f88ef 100644
--- a/mace/ops/concat_benchmark.cc
+++ b/mace/ops/concat_benchmark.cc
@@ -15,7 +15,7 @@ static void ConcatHelper(int iters, int concat_dim, int dim1) {
   OpDefBuilder("Concat", "ConcatBM")
       .Input("Input0")
       .Input("Input1")
-      .Input("Axis")
+      .AddIntArg("axis", concat_dim)
       .Output("Output")
       .Finalize(net.NewOperatorDef());
 
@@ -23,7 +23,6 @@ static void ConcatHelper(int iters, int concat_dim, int dim1) {
   const int kDim0 = 100;
   net.AddRandomInput<DeviceType::CPU, T>("Input0", {kDim0, dim1});
   net.AddRandomInput<DeviceType::CPU, T>("Input1", {kDim0, dim1});
-  net.AddInputFromArray<DeviceType::CPU, int32_t>("Axis", {}, {concat_dim});
 
   // Warm-up
   for (int i = 0; i < 5; ++i) {
@@ -60,14 +59,13 @@ static void OpenclConcatHelper(int iters,
   // Add input data
   net.AddRandomInput<DeviceType::OPENCL, float>("Input0", shape0);
   net.AddRandomInput<DeviceType::OPENCL, float>("Input1", shape1);
-  net.AddInputFromArray<DeviceType::OPENCL, int32_t>("Axis", {}, {concat_dim});
 
   BufferToImage<DeviceType::OPENCL, T>(net, "Input0", "InputImage0", kernels::BufferType::IN_OUT);
   BufferToImage<DeviceType::OPENCL, T>(net, "Input1", "InputImage1", kernels::BufferType::IN_OUT);
   OpDefBuilder("Concat", "ConcatBM")
       .Input("InputImage0")
       .Input("InputImage1")
-      .Input("Axis")
+      .AddIntArg("axis", concat_dim)
       .Output("OutputImage")
       .AddIntArg("T", static_cast<int>(DataTypeToEnum<T>::value))
       .Finalize(net.NewOperatorDef());