 
     switch (GetDesiredBlockDim(feature_size)) {
       FIXED_BLOCK_DIM_CASE(
-          LayerNormForward<T, LayerNormParamType<T>,
+          LayerNormForward<T, U,
                            kBlockDim><<<batch_size, kBlockDim, 0, stream>>>(
               x_data, scale_data, bias_data, y_data, mean_data, var_data,
               epsilon, feature_size));


template <typename T>
__inline__ __device__ T rsqrt(const T val) {
  return ::rsqrt(val);
}

template <>

}

template <>
__inline__ __device__ double rsqrt(const double val) {
  return rsqrt(val);
}

#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
template <>
__inline__ __device__ half rsqrt(const half val) {
  return hrsqrt(val);
}
#endif

template <typename T, typename U, int BlockDim>
__global__ void LayerNormForward(const T *x, const U *scale, const U *bias,
