   int nbatch = data_shape[0];
   int nchannel = data_shape[1];
   float eps = param.eps;
-  std::cout << "nbatch = " << nbatch << ", nchannel = " << nchannel << ", eps = " << eps << std::endl;
   int ngrid_x = (nbatch > kMaxGridDim) ? (nbatch + kBaseGridNum - 1) / kBaseGridNum : nbatch;
   int ngrid_y = (nbatch > kMaxGridDim) ? kBaseGridNum : 1;
-  std::cout << "ngrid = " << ngrid_x << " " << ngrid_y << std::endl;
   int nthread_y = 0;
   const dim3 dimGrid(ngrid_x, ngrid_y, 1);
   if(nchannel <= 32) {

  extern __shared__ char buf[];  // Shared memory size

  if (bid < nbatch) {
    if(threadIdx.x == 0 and threadIdx.y == 0) printf("bid=%d, nbatch=%d\n", bid, nbatch);
    int tid = threadIdx.x + threadIdx.y * blockDim.x;
    const DType* col_vals = in_data + bid * nchannel;
    // Each thread takes charge of 4 consecutive numbers

      DType meanB = WARP_SHFL(mean, src_lane);
      DType sigma2B = WARP_SHFL(sigma2, src_lane);
      DType countB = WARP_SHFL(count, src_lane);
      chan_merge_partition(meanB, sigma2B, countB, mean, sigma2, count);
    }
    if (blockDim.y == 1) {
      mean = WARP_SHFL(mean, 0);
