| 1408 | """ |
| 1409 | |
| 1410 | def __init__(self, config): |
| 1411 | super(BertForQuestionAnswering, self).__init__(config) |
| 1412 | self.bert = BertModel(config) |
| 1413 | # TODO check with Google if it's normal there is no dropout on the token classifier of SQuAD in the TF version |
| 1414 | # self.dropout = nn.Dropout(config.hidden_dropout_prob) |
| 1415 | self.qa_outputs = nn.Linear(config.hidden_size, 2) |
| 1416 | # self.qa_outputs = mpu.RowParallelLinear( |
| 1417 | # input_size=config.hidden_size, |
| 1418 | # output_size=2, |
| 1419 | # bias=True, |
| 1420 | # input_is_parallel=True, |
| 1421 | # stride=1, |
| 1422 | # init_method=normal_init_method(mean=0.0, |
| 1423 | # std=config.initializer_range)) |
| 1424 | self.apply(self.init_bert_weights) |
| 1425 | |
| 1426 | def forward(self, input_ids, token_type_ids=None, attention_mask=None, start_positions=None, end_positions=None, |
| 1427 | checkpoint_activations=False): |