Fix dtypes

candle-nn/src/layer_norm.rs

@@ -31,7 +31,7 @@
 use std::time::{SystemTime, UNIX_EPOCH};
 
 use candle::{backend::BackendStorage, cuda_backend::{cudarc::driver::{LaunchAsync, LaunchConfig}, kernel_name, CudaDType}, DType, Device, Result, Storage, Tensor, D};
-//pub use candle_kernels as kernels;
+pub use candle_kernels as kernels;
 
 #[derive(Debug, Clone, Copy, PartialEq)]
 pub struct LayerNormConfig {
@@ -175,13 +175,13 @@ impl crate::Module for LayerNorm {
             block_dim: (K_CUDABLOCK_REDUCE_NUM_THREADS,1,1),
             shared_mem_bytes: 0,
         };
-        let rowwisemoments = cuda_dev.get_or_load_func(&kernel_name::<T>("rowwisemoments"), kernels::LAYERNORM)?;
+        let rowwisemoments = cuda_dev.get_or_load_func(&format!("rowwisemoments_{}", x.dtype().as_str()), kernels::LAYERNORM)?;
         let params = (n, self.eps, x_storage, mean_storage, rstd_storage);
         unsafe { rowwisemoments.launch(cfg_1, params) };
 
         panic!("Done!");
 
-        let layernorm = cuda_dev.get_or_load_func(&kernel_name::<T>("layernorm"), kernels::LAYERNORM)?;
+        let layernorm = cuda_dev.get_or_load_func(&format!("layernorm_{}", x.dtype().as_str()), kernels::LAYERNORM)?;
         todo!()
     }
 }