| 573 | |
| 574 | |
| 575 | class BertPooler(nn.Module): |
| 576 | def __init__(self, config): |
| 577 | super(BertPooler, self).__init__() |
| 578 | self.dense = nn.Linear(config.hidden_size, config.hidden_size) |
| 579 | self.activation = nn.Tanh() |
| 580 | |
| 581 | def forward(self, hidden_states): |
| 582 | # We "pool" the model by simply taking the hidden state corresponding |
| 583 | # to the first token. |
| 584 | first_token_tensor = hidden_states[:, 0] |
| 585 | pooled_output = self.dense(first_token_tensor) |
| 586 | pooled_output = self.activation(pooled_output) |
| 587 | return pooled_output |
| 588 | |
| 589 | |
| 590 | class BertPredictionHeadTransform(nn.Module): |