| 335 | |
| 336 | |
| 337 | class BertOutput(nn.Module): |
| 338 | def __init__(self, config): |
| 339 | super().__init__() |
| 340 | self.dense = nn.Linear(config.intermediate_size, config.hidden_size) |
| 341 | self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps) |
| 342 | self.dropout = nn.Dropout(config.hidden_dropout_prob) |
| 343 | |
| 344 | def forward(self, hidden_states, input_tensor): |
| 345 | hidden_states = self.dense(hidden_states) |
| 346 | hidden_states = self.dropout(hidden_states) |
| 347 | hidden_states = self.LayerNorm(hidden_states + input_tensor) |
| 348 | return hidden_states |
| 349 | |
| 350 | |
| 351 | class BertLayer(nn.Module): |