Method load

cake-core/src/models/flux/flux2_model.rs:185–188

(dim: usize, eps: f64, vb: VarBuilder)

impl QkNorm {
    fn load(dim: usize, eps: f64, vb: VarBuilder) -> Result<Self> {
        let weight = vb.get(dim, "weight")?;
        Ok(Self { weight, eps })
    }

    fn forward(&self, x: &Tensor) -> Result<Tensor> {
        let in_dtype = x.dtype();