     sum_val1 /= nchannel;
     // 2. Calculate the gradient as
     //      out_grad * gamma / std - sum_val0 - (x - mean) / std * sum_val1
-    for(int l = tid; l < nchannel; l += nthread) {
-      DType ele_out_grad = out_grad[bid * nchannel + l];
-      DType ele_x = in_data[bid * nchannel + l];
-      if(data_addto) {
+    for (int l = tid; l < nchannel; l += nthread) {
+      AType ele_out_grad = static_cast<AType>(out_grad[bid * nchannel + l]);
+      AType ele_x = static_cast<AType>(in_data[bid * nchannel + l]);
+      AType ele_gamma = static_cast<AType>(gamma[l]);
+      if (data_addto) {
         data_grad[bid * nchannel + l] +=
-          ele_out_grad * gamma[l] * invstd_eps - sum_val0
-                                               - (ele_x - mean) * invstd_eps * sum_val1;
+          static_cast<DType>(ele_out_grad * ele_gamma * invstd_eps
+                               - sum_val0 - (ele_x - mean) * invstd_eps * sum_val1);
       } else {
         data_grad[bid * nchannel + l] =
-          ele_out_grad * gamma[l] * invstd_eps - sum_val0
-                                               - (ele_x - mean) * invstd_eps * sum_val1;
+          static_cast<DType>(ele_out_grad * ele_gamma * invstd_eps - sum_val0
+                                               - (ele_x - mean) * invstd_eps * sum_val1);
       }
     }
   }

    // within a warp of threads.
    // After calling the function, threadIdx.x == 0 will store the result of
    // the aggregated (mean, sigma2, counts).
//#pragma unroll
//    for (int l = 0; l <= 4; ++l) {
//      int src_lane = (threadIdx.x + (1<<l)) & 31;
//      DType meanB = WARP_SHFL(mean, src_lane);
//      DType sigma2B = WARP_SHFL(sigma2, src_lane);
//      DType countB = WARP_SHFL(count, src_lane);
//      ChanMergePartition(meanB, sigma2B, countB, mean, sigma2, count);
//    }
    for(int mask = 16; mask > 0; mask >>= 1) {
      DType meanB = WARP_SHFL_XOR(mean, mask);
      DType sigma2B = WARP_SHFL_XOR(sigma2, mask);
      DType countB = WARP_SHFL_XOR(count, mask);
      ChanMergePartition(meanB, sigma2B, countB, mean, sigma2, count);
    }
    if (blockDim.y > 1) {
      // Inter-warp reduction. Copy the upper-half of the warps to shared memory
      // and merge with the lower-half warp
      DType* mean_buf = reinterpret_cast<DType*>(buf);
      DType* sigma2_buf = reinterpret_cast<DType*>(buf + sizeof(DType) * blockDim.y / 2 * 32);
      DType* count_buf = reinterpret_cast<DType*>(buf + sizeof(DType) * blockDim.y * 32);
      for (int offset = blockDim.y / 2; offset > 0; offset >>= 1) {
        if (threadIdx.y >= offset && threadIdx.y < 2 * offset) {
          const int idx = (threadIdx.y - offset) * blockDim.x + threadIdx.x;

 *
 *  TODO(sxjscience) Explore the possibility of int lhs_count and rhs_count
 */
template<typename DType, typename IType>
__device__ __inline__ void ChanMergePartition(const DType lhs_mean,
                                              const DType lhs_sigma2,
                                              const IType lhs_count,
                                              DType& rhs_mean,         //NOLINT
                                              DType& rhs_sigma2,       //NOLINT
                                              IType& rhs_count) {      //NOLINT
  DType delta = rhs_mean - lhs_mean;
  DType nA = static_cast<DType>(lhs_count);
  DType nB = static_cast<DType>(rhs_count);
  rhs_count = nA + nB;
  if (rhs_count > DType(0)) {
    nA = nA / rhs_count;
