| 1345 | self.apply(self.init_bert_weights) |
| 1346 | |
| 1347 | def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None, checkpoint_activations=False): |
| 1348 | sequence_output, _ = self.bert(input_ids, token_type_ids, attention_mask, output_all_encoded_layers=False, |
| 1349 | checkpoint_activations=checkpoint_activations) |
| 1350 | with mpu.get_cuda_rng_tracker().fork(): |
| 1351 | sequence_output = self.dropout(sequence_output) |
| 1352 | logits = self.classifier(sequence_output) |
| 1353 | |
| 1354 | if labels is not None: |
| 1355 | loss_fct = CrossEntropyLoss() |
| 1356 | loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) |
| 1357 | return loss |
| 1358 | else: |
| 1359 | return logits |
| 1360 | |
| 1361 | |
| 1362 | class BertForQuestionAnswering(PreTrainedBertModel): |