diff --git a/mace/kernels/opencl/cl/space_to_batch.cl b/mace/kernels/opencl/cl/space_to_batch.cl
index 24a07712052e1ea995b5e804ef00f17c7b1296dc..921a3bf82e402439198081884a9bfefb164b4161 100644
--- a/mace/kernels/opencl/cl/space_to_batch.cl
+++ b/mace/kernels/opencl/cl/space_to_batch.cl
@@ -1,6 +1,6 @@
-void kernel space_to_batch(global float* space_data_ptr,
-                           global const int* block_shape_ptr,
-                           global const int* paddings_ptr,
+void kernel space_to_batch(global float *space_data_ptr,
+                           global const int *block_shape_ptr,
+                           global const int *paddings_ptr,
                            private const int space_batch,
                            private const int space_channel,
                            private const int space_height,
@@ -28,10 +28,12 @@ void kernel space_to_batch(global float* space_data_ptr,
                                batch_pixel_height_idx * block_height;
   int space_pixel_width_idx = (remaining_batch_idx % block_width) +
                               batch_pixel_width_idx * block_width;
+
   const int batch_data_offset = batch_idx * (space_channel * batch_height * batch_width) +
                                 (batch_channel_idx * batch_height * batch_width) +
                                 batch_pixel_height_idx * batch_width +
                                 batch_pixel_width_idx;
+
   space_pixel_height_idx -= padding_height_start;
   space_pixel_width_idx -= padding_width_start;
   const int space_data_offset = space_idx * (space_channel * space_height * space_width) +
diff --git a/mace/ops/batch_to_space.h b/mace/ops/batch_to_space.h
index d1080f37fd7d6c4bf4c6b16916a001494e7e81a9..14b6444553860935d4fe4add7e353727b0d74c96 100644
--- a/mace/ops/batch_to_space.h
+++ b/mace/ops/batch_to_space.h
@@ -22,6 +22,7 @@ static void BatchToSpaceHelper(const Tensor *input_tensor,
 
   const index_t block_dims = block_shape_tensor->dim(0);
   MACE_CHECK(block_dims == cropped_tensor->dim(0) && 2 == cropped_tensor->dim(1));
+  // TODO change tensor to attribute if needed based on the benchmark
   Tensor::MappingGuard block_shape_tensor_mapper(block_shape_tensor);
   Tensor::MappingGuard cropped_tensor_mapper(cropped_tensor);
   const int *block_shape_ptr = block_shape_tensor->data<int>();