 template<typename DType>
 __global__ void LayerNormFusedForwardKernelContig(const int nbatch,
                                                   const int nchannel,
-                                                  const float eps,
+                                                  const DType eps,
                                                   const DType* __restrict__ in_data,
                                                   const DType* __restrict__ gamma,
                                                   const DType* __restrict__ beta,

    int nshared = nthread_y > 1 ? nthread_y * sizeof(DType) + (nthread_y / 2) * sizeof(DType) : 0;
    CheckLaunchParam(dimGrid, dimBlock);
    LayerNormFusedForwardKernelContig<<<dimGrid, dimBlock, nshared, stream>>>
     (nbatch, nchannel, eps,
      in_data.dptr<DType>(), gamma.dptr<DType>(), beta.dptr<DType>(),
      out_data.dptr<DType>(), mean_data.dptr<DType>(), std_data.dptr<DType>());
    MSHADOW_CUDA_POST_KERNEL_CHECK(LayerNormFusedForwardKernelContig);

    // Write the out_data and var_data
    if(threadIdx.x == 0 && threadIdx.y == 0) {
      mean_data[bid] = mean;
      std_data[bid] = std_eps;
    }
  }
}
