| 660 | |
| 661 | |
| 662 | class BertOnlyNSPHead(nn.Module): |
| 663 | def __init__(self, config): |
| 664 | super(BertOnlyNSPHead, self).__init__() |
| 665 | self.seq_relationship = nn.Linear(config.hidden_size, 2) |
| 666 | |
| 667 | def forward(self, pooled_output): |
| 668 | seq_relationship_score = self.seq_relationship(pooled_output) |
| 669 | return seq_relationship_score |
| 670 | |
| 671 | |
| 672 | class BertPreTrainingHeads(nn.Module): |