| 510 | |
| 511 | |
| 512 | class BertLayer(nn.Module): |
| 513 | def __init__(self, config): |
| 514 | super(BertLayer, self).__init__() |
| 515 | self.attention = BertAttention(config) |
| 516 | self.intermediate = BertIntermediate(config) |
| 517 | self.output = BertOutput(config) |
| 518 | |
| 519 | def forward(self, hidden_states, attention_mask): |
| 520 | attention_output = self.attention(hidden_states, attention_mask) |
| 521 | intermediate_output = self.intermediate(attention_output) |
| 522 | layer_output = self.output(intermediate_output, attention_output) |
| 523 | return layer_output |
| 524 | |
| 525 | |
| 526 | class BertEncoder(nn.Module): |