Do not implement Module for BatchNorm. (#1513)

This commit is contained in:
Laurent Mazare
2024-01-01 10:13:13 +01:00
committed by GitHub
parent 1fb2dd905c
commit b0fe5e4453
9 changed files with 31 additions and 33 deletions

View File

@ -147,7 +147,7 @@ fn conv(vb: VarBuilder, index: usize, p: usize, b: &Block) -> Result<(usize, Bl)
let func = candle_nn::func(move |xs| {
let xs = conv.forward(xs)?;
let xs = match &bn {
Some(bn) => bn.forward(&xs)?,
Some(bn) => xs.apply_t(bn, false)?,
None => xs,
};
let xs = if leaky {