| 106 | |
| 107 | impl DecoderBlock { |
| 108 | fn load(vb: VarBuilder, channels: usize, eps: f64, backend: Arc<dyn ComputeBackend>) -> Result<Self> { |
| 109 | let norm_weight = vb.pp("norm").get(channels, "weight")?; |
| 110 | let gamma = vb.get(channels, "gamma")?; |
| 111 | |
| 112 | let conv_vb = vb.pp("mixer").pp("conv").pp("conv").pp("conv"); |
| 113 | let mixer_weight = conv_vb.get((channels, 1, 7), "weight")?.squeeze(1)?; |
| 114 | let mixer_bias = conv_vb.get(channels, "bias")?; |
| 115 | |
| 116 | let ffn_norm_weight = vb.pp("ffn_norm").get(channels, "weight")?; |
| 117 | let ffn_gamma = vb.get(channels, "ffn_gamma")?; |
| 118 | let ffn_linear1_weight = vb.pp("ffn").pp("linear1").get((channels * 4, channels), "weight")?; |
| 119 | let ffn_linear1_bias = vb.pp("ffn").pp("linear1").get(channels * 4, "bias").ok(); |
| 120 | let ffn_linear2_weight = vb.pp("ffn").pp("linear2").get((channels, channels * 4), "weight")?; |
| 121 | let ffn_linear2_bias = vb.pp("ffn").pp("linear2").get(channels, "bias").ok(); |
| 122 | |
| 123 | Ok(Self { |
| 124 | norm_weight, |
| 125 | ffn_norm_weight, |
| 126 | eps: eps as f32, |
| 127 | gamma, |
| 128 | mixer_weight, |
| 129 | mixer_bias, |
| 130 | ffn_gamma, |
| 131 | ffn_linear1_weight, |
| 132 | ffn_linear1_bias, |
| 133 | ffn_linear2_weight, |
| 134 | ffn_linear2_bias, |
| 135 | backend, |
| 136 | }) |
| 137 | } |
| 138 | |
| 139 | fn forward(&self, x: &Tensor) -> Result<Tensor> { |
| 140 | let channels = x.dim(1)?; |