 void GetGammaBetaGradKernelParams(const int nbatch, const int nchannel,
                                   dim3* part_grad_block_dim, dim3* part_grad_grid_dim,
                                   dim3* gb_block_dim, dim3* gb_grid_dim,
-                                  int* row_repeat, int* npart) {
+                                  int* npart) {
   *npart = 16;
   *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);
   *gb_grid_dim = dim3((nchannel + 32 - 1) / 32);
   CheckLaunchParam(*part_grad_grid_dim, *part_grad_block_dim);

  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 * row_stride;
  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) {
      for(int i = 0; i < row_repeat; ++i) {
        int r_offset = i * blockDim.y + threadIdx.y;
        int r = r_b + r_offset;
        if(r < r_end) {
          DType local_mean = mean_data[r];
          DType local_std = std_data[r];
          int read_idx = r * nchannel + c;
          local_gamma_grad += (in_data[read_idx] - local_mean) / local_std
                                                               * out_grad[read_idx];
          local_beta_grad += out_grad[read_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 > 1;  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];

      __syncthreads();
      mean = mean_buf[threadIdx.x];
      sigma2 = sigma2_buf[threadIdx.x] / nchannel;
    } else {
      sigma2 /= nchannel;
    }
    // Calculate the out_data: gamma * (x - mean) / sqrt(var + eps) + beta
    AType std_eps = sqrt(sigma2 + eps);
    AType invstd_eps = DType(1.0) / std_eps;
    DType* out_col_val = out_data + bid * nchannel;

    if (gamma != NULL && beta != NULL) {
      for (int i = tid; i < nchannel; i += nthread) {
        out_col_val[i] = gamma[i] * static_cast<DType>(invstd_eps *
                                                       (static_cast<AType>(col_vals[i]) - mean))
                                                         + beta[i];
      }
    } else if (gamma == NULL && beta != NULL) {
      for (int i = tid; i < nchannel; i += nthread) {
        out_col_val[i] = static_cast<DType>(invstd_eps * (static_cast<AType>(col_vals[i]) - mean))
                                                       + beta[i];
      }
    } else if (gamma != NULL && beta == NULL) {
      for (int i = tid; i < nchannel; i += nthread) {
        out_col_val[i] = gamma[i] * static_cast<DType>(invstd_eps *
                                                       (static_cast<AType>(col_vals[i]) - mean));
      }
    } else {
      for (int i = tid; i < nchannel; i += nthread) {
        out_col_val[i] = static_cast<DType>(invstd_eps * (static_cast<AType>(col_vals[i]) - mean));
      }
    }
    // Write the out_data and var_data
    if (threadIdx.x == 0 && threadIdx.y == 0) {
      mean_data[bid] = static_cast<DType>(mean);
      std_data[bid] = static_cast<DType>(std_eps);
    }
  }
}
