   DType* d_buf = reinterpret_cast<DType*>(buf);
   const int npart = gridDim.y;
   const int block_row_num = (nbatch + npart - 1) / npart;
+  const int nthread = blockDim.x * blockDim.y;
+  const int tid = threadIdx.y * blockDim.x + threadIdx.x;
   // The rows are divided into `npart` parts. Each threadblock calculates the reduction result
   // within the corresponding row ranges.
   const int c = blockIdx.x * blockDim.x + threadIdx.x;
   int r_begin = blockIdx.y * block_row_num;
   int r_end = min((blockIdx.y + 1) * block_row_num, nbatch);
-  const int row_stride = blockDim.x + 1;
   DType* buf_gamma_grad = d_buf;
-  DType* buf_beta_grad = d_buf + blockDim.y * row_stride * row_repeat;
-  for(int i = 0; i < row_repeat; i++) {
-    int idx = (i * blockDim.y + threadIdx.y) * row_stride + threadIdx.x;
-    buf_gamma_grad[idx] = 0;
-    buf_beta_grad[idx] = 0;
-  }
+  DType* buf_beta_grad = d_buf + blockDim.y * blockDim.x;
+  DType local_gamma_grad = 0;
+  DType local_beta_grad = 0;
 
   if(c < nchannel) {
     for(int r_b = r_begin; r_b < r_end; r_b += blockDim.y * row_repeat) {

                                  dim3* gb_block_dim, dim3* gb_grid_dim,
                                  int* row_repeat, int* npart) {
  *npart = 16;
  *part_grad_block_dim = dim3(32, 4);
  *part_grad_grid_dim = dim3((nchannel + 32 - 1) / 32, *npart);
  *row_repeat = 4;
  *gb_block_dim = dim3(32, *npart);

        ctx.requested[0].get_space_typed<gpu, 1, DType>(Shape1(2 * npart * nchannel), s);
      DType* part_gamma_grad_ptr = workspace.dptr_;
      DType* part_beta_grad_ptr = workspace.dptr_ + npart * nchannel;
      const int nshared_K1 = 2 * part_grad_block_dim.x * part_grad_block_dim.y * sizeof(DType);
      const int nshared_K2 = 2 * gb_block_dim.x * gb_block_dim.y * sizeof(DType);
      DType* gamma_grad_ptr = (gamma_grad_req != kNullOp) ? gamma_grad.dptr<DType>() : nullptr;
      DType* beta_grad_ptr = (beta_grad_req != kNullOp) ? beta_grad.dptr<DType>() : nullptr;
