| 446 | |
| 447 | |
| 448 | class BertLayer(nn.Module): |
| 449 | def __init__(self, config): |
| 450 | super(BertLayer, self).__init__() |
| 451 | self.attention = BertAttention(config) |
| 452 | self.intermediate = BertIntermediate(config) |
| 453 | self.output = BertOutput(config) |
| 454 | |
| 455 | def forward(self, hidden_states, attention_mask): |
| 456 | attention_output = self.attention(hidden_states, attention_mask) |
| 457 | intermediate_output = self.intermediate(attention_output) |
| 458 | layer_output = self.output(intermediate_output, attention_output) |
| 459 | return layer_output |
| 460 | |
| 461 | class BertEncoder(nn.Module): |
| 462 | def __init__(self, config): |