Method load

cake-core/src/models/flux/flux2_vae.rs:23–27

(in_ch: usize, out_ch: usize, kernel: usize, padding: usize, vb: VarBuilder)

impl RawConv2d {
    fn load(in_ch: usize, out_ch: usize, kernel: usize, padding: usize, vb: VarBuilder) -> Result<Self> {
        let weight = vb.get((out_ch, in_ch, kernel, kernel), "weight")?;
        let bias = vb.get(out_ch, "bias").ok();
        Ok(Self { weight, bias, padding })
    }

    fn forward(&self, x: &Tensor, backend: &dyn ComputeBackend) -> Result<Tensor> {
        backend.conv2d(x, &self.weight, self.bias.as_ref(), self.padding, 1, 1, 1)