| 508 | |
| 509 | |
| 510 | class BertOnlyMLMHead(nn.Module): |
| 511 | def __init__(self, config): |
| 512 | super().__init__() |
| 513 | self.predictions = BertLMPredictionHead(config) |
| 514 | |
| 515 | def forward(self, sequence_output): |
| 516 | prediction_scores = self.predictions(sequence_output) |
| 517 | return prediction_scores |
| 518 | |
| 519 | |
| 520 | class BertOnlyNSPHead(nn.Module): |