| 518 | |
| 519 | |
| 520 | class BertOnlyNSPHead(nn.Module): |
| 521 | def __init__(self, config): |
| 522 | super().__init__() |
| 523 | self.seq_relationship = nn.Linear(config.hidden_size, 2) |
| 524 | |
| 525 | def forward(self, pooled_output): |
| 526 | seq_relationship_score = self.seq_relationship(pooled_output) |
| 527 | return seq_relationship_score |
| 528 | |
| 529 | |
| 530 | class BertPreTrainingHeads(nn.Module): |