  *  TODO(sxjscience) Explore the possibility of int lhs_count and rhs_count
  */
 template<typename DType>
-__device__ __forceinline__ void chan_merge_partition(const DType lhs_mean,
+__device__ void chan_merge_partition(const DType lhs_mean,
                                      const DType lhs_sigma2,
                                      const DType lhs_count,
                                      DType& rhs_mean,

      }
      chan_merge_partition(meanB, sigma2B, countB, mean, sigma2, count);
    }
     */
    warp_merge_mean_sigma2(mean, sigma2, count);
    if(bid == 0 and threadIdx.x == 0) {
      printf("threadIdx.y = %d, mean = %g, sigma2 = %g, count = %d\n", threadIdx.y, mean, sigma2, count);
    }

  }
}

/* Fused CUDA kernel for layer normalization. It computes the LayerNorm when axis=-1.
 * Shape of the input tensors:
 *      in_data = (nbatch, nchannel)
