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

                                                  DType* out_data,
                                                  DType* mean_data,
                                                  DType* std_data) {
  int bid = blockIdx.x + blockIdx.y * gridDim.x;
  const int nthread = blockDim.x * blockDim.y;
  DType count = 0;
  DType mean = 0;

      for (int i = tid; i < nchannel; i += nthread) {
        out_col_val[i] = invstd_eps * (col_vals[i] - mean);
      }
    }
    // Write the out_data and var_data
    if(threadIdx.x == 0 && threadIdx.y == 0) {
