VarBuilder cleanup (#627)

* VarBuilder cleanup.

* Implement the basic varbuilders.

* Add the sharded code.

* Proper support for tensor sharding.
This commit is contained in:
Laurent Mazare
2023-08-27 18:03:26 +01:00
committed by GitHub
parent be471d50ab
commit 4c338b0cd9
12 changed files with 409 additions and 291 deletions

View File

@ -79,7 +79,7 @@ pub fn group_norm(
eps: f64,
vb: crate::VarBuilder,
) -> Result<GroupNorm> {
let weight = vb.get_or_init(num_channels, "weight", crate::Init::Const(1.))?;
let bias = vb.get_or_init(num_channels, "bias", crate::Init::Const(0.))?;
let weight = vb.get_with_hints(num_channels, "weight", crate::Init::Const(1.))?;
let bias = vb.get_with_hints(num_channels, "bias", crate::Init::Const(0.))?;
GroupNorm::new(weight, bias, num_channels, num_groups, eps)
}