| 34 | |
| 35 | impl EncoderBlock { |
| 36 | fn load(vb: VarBuilder, channels: usize, eps: f64, backend: Arc<dyn ComputeBackend>) -> Result<Self> { |
| 37 | let norm_weight = vb.pp("norm").get(channels, "weight")?; |
| 38 | let gamma = vb.get(channels, "gamma")?; |
| 39 | |
| 40 | let conv_vb = vb.pp("mixer").pp("conv").pp("conv").pp("conv"); |
| 41 | let mixer_weight = conv_vb.get((channels, 1, 7), "weight")?.squeeze(1)?; |
| 42 | let mixer_bias = conv_vb.get(channels, "bias")?; |
| 43 | |
| 44 | let ffn_norm_weight = vb.pp("ffn_norm").get(channels, "weight")?; |
| 45 | let ffn_gamma = vb.get(channels, "ffn_gamma")?; |
| 46 | let ffn_linear1_weight = vb.pp("ffn").pp("linear1").get((channels * 4, channels), "weight")?; |
| 47 | let ffn_linear1_bias = vb.pp("ffn").pp("linear1").get(channels * 4, "bias").ok(); |
| 48 | let ffn_linear2_weight = vb.pp("ffn").pp("linear2").get((channels, channels * 4), "weight")?; |
| 49 | let ffn_linear2_bias = vb.pp("ffn").pp("linear2").get(channels, "bias").ok(); |
| 50 | |
| 51 | Ok(Self { |
| 52 | norm_weight, ffn_norm_weight, eps: eps as f32, |
| 53 | gamma, mixer_weight, mixer_bias, ffn_gamma, |
| 54 | ffn_linear1_weight, ffn_linear1_bias, ffn_linear2_weight, ffn_linear2_bias, |
| 55 | backend, |
| 56 | }) |
| 57 | } |
| 58 | |
| 59 | fn forward(&self, x: &Tensor) -> Result<Tensor> { |
| 60 | let channels = x.dim(1)?; |