Dynamically dispatch on dtype

EricLBuehler · Mar 12, 2024 · 37f5336 · 37f5336
1 parent 53faff9
commit 37f5336
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/candle-nn/src/layer_norm.rs b/candle-nn/src/layer_norm.rs
@@ -194,7 +194,7 @@ impl crate::Module for LayerNorm {
             eps: f32,
         }
 
-        fn run<T: WithDType>(dev: CudaDevice, storage: &candle::CudaStorage, layout: &Layout) -> Result<CudaStorage> {
+        fn run<T: WithDType + CudaDType>(dev: CudaDevice, storage: &candle::CudaStorage, layout: &Layout) -> Result<CudaStorage> {
             let slice = storage.as_cuda_slice::<T>()?;
             let slice = match layout.contiguous_offsets() {
                 None => candle::bail!("input has to be contiguous"),