Fix batch size

EricLBuehler · Mar 12, 2024 · 7bcf36a · 7bcf36a
1 parent 89d21ce
commit 7bcf36a
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
@@ -156,7 +156,7 @@ impl crate::Module for LayerNorm {
             Device::Cpu | Device::Metal(_) => return self.forward_slow(x),
             Device::Cuda(dev) => dev
         };
-        let (m,n) = x.dims2()?;
+        let (b, m,n) = x.dims3()?;
 
         let mean = x.zeros_like()?;
         let rstd = x.zeros_like()?;