| 164 | |
| 165 | |
| 166 | class TransformerEncoderLayer(nn.Module): |
| 167 | def __init__(self, config: T5Config): |
| 168 | super().__init__() |
| 169 | self.attention = MultiHeadAttention(config) |
| 170 | self.ln1 = nn.RMSNorm(config.d_model, eps=config.layer_norm_epsilon) |
| 171 | self.ln2 = nn.RMSNorm(config.d_model, eps=config.layer_norm_epsilon) |
| 172 | self.dense = DenseActivation(config) |
| 173 | |
| 174 | def __call__(self, x, mask): |
| 175 | y = self.ln1(x) |
| 176 | y, _ = self.attention(y, y, y, mask=mask) |
| 177 | x = x + y |
| 178 | |
| 179 | y = self.ln2(x) |
| 180 | y = self.dense(y) |
| 181 | return x + y |
| 182 | |
| 183 | |
| 184 | class TransformerEncoder(nn.Module): |