huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 10:38:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Expose the fields from batch-norm. (#1176)

Browse Source
This commit is contained in:
Laurent Mazare
2023-10-25 15:35:32 +01:00
committed by GitHub
parent c698e17619
commit 0acd16751d
1 changed files with 12 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
candle-nn/src/batch_norm.rs
View File

@ -100,9 +100,19 @@ impl BatchNorm {
num_features, num_features,
}) })
} }
}
impl BatchNorm { pub fn running_mean(&self) -> &Tensor {
&self.running_mean
}
pub fn running_var(&self) -> &Tensor {
&self.running_var
}
pub fn weight_and_bias(&self) -> Option<(&Tensor, &Tensor)> {
self.weight_and_bias.as_ref().map(|v| (&v.0, &v.1))
}
pub fn forward_learning(&self, x: &Tensor) -> Result<Tensor> { pub fn forward_learning(&self, x: &Tensor) -> Result<Tensor> {
let x_dtype = x.dtype(); let x_dtype = x.dtype();
let internal_dtype = match x_dtype { let internal_dtype = match x_dtype {