huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-21 12:20:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Do not implement Module for BatchNorm. (#1513)

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

6
candle-transformers/src/models/convmixer.rs
View File

@ -40,8 +40,8 @@ fn block(dim: usize, kernel_size: usize, vb: VarBuilder) -> Result<impl Module>
let conv2 = candle_nn::conv2d(dim, dim, 1, Default::default(), vb.pp(1))?;
let bn2 = batch_norm(dim, 1e-5, vb.pp(3))?;
Ok(candle_nn::func(move |xs| {
let ys = xs.apply(&conv1)?.gelu_erf()?.apply(&bn1)?;
(xs + ys)?.apply(&conv2)?.gelu_erf()?.apply(&bn2)
let ys = xs.apply(&conv1)?.gelu_erf()?.apply_t(&bn1, false)?;
(xs + ys)?.apply(&conv2)?.gelu_erf()?.apply_t(&bn2, false)
}))
}
@ -64,7 +64,7 @@ fn convmixer(
.collect::<Result<Vec<_>>>()?;
let fc = candle_nn::linear(dim, nclasses, vb.pp(25))?;
Ok(candle_nn::func(move |xs| {
let mut xs = xs.apply(&conv1)?.gelu_erf()?.apply(&bn1)?;
let mut xs = xs.apply(&conv1)?.gelu_erf()?.apply_t(&bn1, false)?;
for block in blocks.iter() {
xs = xs.apply(block)?
}