   int inC = data.shape_[3];
   int outC = weight.shape_[0];
   const int grid_dim_x = B * H * W;
-  const int grid_dim_y = (outC + TILE_SIZE - 1) / TILE_SIZE;
+  // const int grid_dim_y = (outC + TILE_SIZE - 1) / TILE_SIZE;
   CHECK_LT(L * K, 128);
   dim3 dimGrid(grid_dim_x, 1);
   dim3 dimBlock(TILE_SIZE, TILE_SIZE);

          weight_shared[ty][tx] = 0.0f;
        }
        // Load the local connection data into shared memory. TODO, accelerate this part using Parallel Reduce.
        if (ic < inC && ty == 0) {
          data_shared[tx] = 0.0f;
          for (int k = 0; k < K; ++k) {
            if (local_connection_ind[l * K + k] >= 0 && local_connection_ind[l * K + k] < H * W) {
              int address = ADDRESS_3D_BHWC(b, local_connection_ind[l * K + k], ic, H * W, inC);
              data_shared[tx] += local_connection_val[l * K + k] * data[address];
            } else {
              data_shared[tx] += local_connection_val[l * K + k] * pad_val;
            }
          }
        }

    int w = index % W;
    int h = (index / W) % H;
    int b = index / W / H;
    // Load the initial data + the local connection values and indices
    if (tid < L * K) {
      int address = ADDRESS_4D_BCHW(b, tid, h, w, L * K, H, W);
      local_connection_val[tid] = values[address];
      local_connection_ind[tid] = __float2int_rn(indices[address]);
    }
    __syncthreads();
#pragma unroll
    for (int oc_base = 0; oc_base < outC; oc_base += TILE_SIZE) {
      if (ty == 0 && oc_base + tx < outC) {
        out_shared[tx] = bias[oc_base + tx];
      }
      __syncthreads();
      int oc = oc_base + ty;
#pragma unroll
      for (int icl = 0; icl < inC * L; icl += TILE_SIZE) {
