diff --git a/mace/BUILD b/mace/BUILD
index 17a4be41ce7a07fb8d8a69f6e9b021b6eb0fa45e..1b95aae048469510fbe8c5d272602519689408e7 100644
--- a/mace/BUILD
+++ b/mace/BUILD
@@ -10,7 +10,7 @@ config_setting(
     name = "android_armv7",
     values = {
         "crosstool_top": "//external:android/crosstool",
-        "android_cpu": "armeabi-v7a",
+        "cpu": "armeabi-v7a",
     },
     visibility = ["//visibility:public"],
 )
@@ -19,7 +19,7 @@ config_setting(
     name = "android_arm64",
     values = {
         "crosstool_top": "//external:android/crosstool",
-        "android_cpu": "arm64-v8a",
+        "cpu": "arm64-v8a",
     },
     visibility = ["//visibility:public"],
 )
diff --git a/mace/core/opencl_allocator.cc b/mace/core/opencl_allocator.cc
index 5acd290eb92e800b8eb5703fa0d3351bf08e47b3..3b393542281266a4564767e732ea703c4371e738 100644
--- a/mace/core/opencl_allocator.cc
+++ b/mace/core/opencl_allocator.cc
@@ -70,7 +70,7 @@ void OpenCLAllocator::Delete(void *buffer) {
 
 void OpenCLAllocator::DeleteImage(void *buffer) {
   if (buffer != nullptr) {
-    cl::Image3D *cl_image = static_cast<cl::Image3D *>(buffer);
+    cl::Image2D *cl_image = static_cast<cl::Image2D *>(buffer);
     delete cl_image;
   }
 }
diff --git a/mace/core/opencl_allocator.h b/mace/core/opencl_allocator.h
index a3758e85c779169080e40df3b4fb643fbb436c27..6aa8fbf93a3d7105b217469a8f0a7c0af5c4d365 100644
--- a/mace/core/opencl_allocator.h
+++ b/mace/core/opencl_allocator.h
@@ -18,8 +18,7 @@ class OpenCLAllocator : public Allocator {
   void *New(size_t nbytes) override;
 
   /*
-   * Only support shape.size() > 1 and collapse first n-2 dimensions to depth.
-   * Use Image3D with RGBA (128-bit) format to represent the image.
+   * Use Image2D with RGBA (128-bit) format to represent the image.
    *
    * @ shape : [depth, ..., height, width ].
    */
diff --git a/mace/kernels/BUILD b/mace/kernels/BUILD
index bd79ac2bbadc33beb45849300bcc82cda89a8723..4352f2a66cd413819e19bf98e70840cbea404d3f 100644
--- a/mace/kernels/BUILD
+++ b/mace/kernels/BUILD
@@ -13,7 +13,7 @@ load("//mace:mace.bzl", "if_android_arm64")
 cc_library(
     name = "kernels",
     srcs = glob(["*.cc"]) + if_android(glob(["opencl/*.cc"])) + if_android_arm64(glob(["neon/*.cc"])),
-    hdrs = glob(["*.h"]) + if_android(glob(["opencl/*.cc"])) + if_android_arm64(glob(["neon/*.cc"])),
+    hdrs = glob(["*.h"]) + if_android(glob(["opencl/*.h"])) + if_android_arm64(glob(["neon/*.h"])),
     copts = [
         "-std=c++11",
         "-fopenmp",
diff --git a/mace/kernels/opencl/buffer_to_image.cc b/mace/kernels/opencl/buffer_to_image.cc
index db44b451f09d5fac68ae750ba8d16f81ae00eebb..be76710a2cea374cd733df19bc24000625a8c4e0 100644
--- a/mace/kernels/opencl/buffer_to_image.cc
+++ b/mace/kernels/opencl/buffer_to_image.cc
@@ -39,13 +39,12 @@ void BufferToImageFunctor<DeviceType::OPENCL, T>::operator()(Tensor *buffer,
       kernel_name = i2b_ ? "arg_image_to_buffer" : "arg_buffer_to_image";
       break;
   }
-  VLOG(0) << kernel_name;
   auto b2f_kernel = runtime->BuildKernel("buffer_to_image",
                                          kernel_name,
                                          built_options);
 
   uint32_t idx = 0;
-  b2f_kernel.setArg(idx++, *(static_cast<const cl::Image3D *>(buffer->buffer())));
+  b2f_kernel.setArg(idx++, *(static_cast<const cl::Image2D *>(buffer->buffer())));
   if (type == ARGUMENT) {
     b2f_kernel.setArg(idx++, static_cast<uint32_t>(buffer->dim(0)));
   } else {
@@ -53,7 +52,7 @@ void BufferToImageFunctor<DeviceType::OPENCL, T>::operator()(Tensor *buffer,
     b2f_kernel.setArg(idx++, static_cast<uint32_t>(buffer->dim(2)));
     b2f_kernel.setArg(idx++, static_cast<uint32_t>(buffer->dim(3)));
   }
-  b2f_kernel.setArg(idx++, *(static_cast<cl::Image3D *>(image->buffer())));
+  b2f_kernel.setArg(idx++, *(static_cast<cl::Image2D *>(image->buffer())));
 
   const size_t gws[3] = {image_shape[0],
                          image_shape[1],