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

          DType local_mean = mean_data[r];
          DType local_std = std_data[r];
          int read_idx = r * nchannel + c;
          int write_idx = r_offset * row_stride + threadIdx.x;
          buf_gamma_grad[write_idx] += (in_data[read_idx] - local_mean) / local_std
                                                                        * out_grad[read_idx];
          buf_beta_grad[write_idx] += out_grad[read_idx];
        }
      }
    }
  }
  __syncthreads();
  DType local_gamma_grad = 0;
  DType local_beta_grad = 0;
  for(int i = 0; i < row_repeat; ++i) {
    int idx = (i * blockDim.y + threadIdx.y) * row_stride + threadIdx.x;
    local_gamma_grad += buf_gamma_grad[idx];
    local_beta_grad += buf_beta_grad[idx];
  }
  buf_gamma_grad[threadIdx.y * row_stride + threadIdx.x] = local_gamma_grad;
  buf_beta_grad[threadIdx.y * row_stride + threadIdx.x] = local_beta_grad;
  __syncthreads();
  for(int offset = blockDim.y/2;  offset > 0;  offset >>= 1) {
    if(threadIdx.y < offset) {
      int idx1 = threadIdx.y * row_stride + threadIdx.x;
      int idx2 = (threadIdx.y + offset) * row_stride + threadIdx.x;
      buf_gamma_grad[idx1] += buf_gamma_grad[idx2];
      buf_beta_grad[idx1] += buf_beta_grad[idx2];
    }

  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);
  DType* buf_gamma_grad = d_buf;
  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) {
