Use custom impl and time it

EricLBuehler · Mar 12, 2024 · 8d9218d · 8d9218d
1 parent 71bd43a
commit 8d9218d
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
@@ -233,7 +233,7 @@ impl crate::Module for LayerNorm {
                 };
                 let elem_count = layout.shape().elem_count();
                 let dst = unsafe { dev.alloc::<f32>(elem_count) }.w()?;
-                let func = dev.get_or_load_func("rms_f32", cuda_kernels::LAYERNORM_KERNELS)?;
+                let func = dev.get_or_load_func("rms_f32", kernels::LAYERNORM_KERNELS)?;
                 let params = (&dst, &slice, self.eps, d1, d2);
                 let cfg = LaunchConfig {
                     grid_dim: (d1, 1, 1),