Try again

candle-nn/src/layer_norm.rs

@@ -34,7 +34,7 @@ pub use candle::cuda_backend::kernels;
 use candle::{
     backend::BackendStorage,
     cuda_backend::{
-        cudarc::driver::{DeviceRepr, LaunchAsync, LaunchConfig},
+        cudarc::driver::{DevicePtr, DeviceRepr, LaunchAsync, LaunchConfig},
         kernel_name, CudaDType, WrapErr,
     },
     CudaDevice, CudaStorage, DType, Device, Result, Storage, Tensor, WithDType, D,
@@ -159,30 +159,30 @@ impl crate::Module for LayerNorm {
         let (m,n) = x.dims2()?;
 
         let x_storage = match &*x.storage_and_layout().0 {
-            Storage::Cuda(s) => s,
+            Storage::Cuda(s) => s.as_cuda_slice()?.device_ptr(),
             _ => unreachable!(),
-        }.slice;
+        };
 
         let mean = x.zeros_like()?;
         let rstd = x.zeros_like()?;
 
         let mean_storage = match &*mean.storage_and_layout().0 {
-            Storage::Cuda(s) => s,
+            Storage::Cuda(s) => s.as_cuda_slice()?.device_ptr(),
             _ => unreachable!(),
-        }.slice;
+        };
 
         let rstd_storage = match &*rstd.storage_and_layout().0 {
-            Storage::Cuda(s) => s,
+            Storage::Cuda(s) => s.as_cuda_slice()?.device_ptr(),
             _ => unreachable!(),
-        }.slice;
+        };
 
         let cfg_1 = LaunchConfig {
             grid_dim: (m as u32,1,1),
             block_dim: (K_CUDABLOCK_REDUCE_NUM_THREADS,1,1),
             shared_mem_bytes: 0,
         };
         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);
+        let params = (n, self.eps, *x_storage, *mean_storage, &rstd_storage);
         unsafe { rowwisemoments.launch(cfg_1, params) };
 
         panic!("Done!");