diff --git a/pytorch_transformers/modeling_bert.py b/pytorch_transformers/modeling_bert.py index 0b1cfd6a..5d3c1606 100644 --- a/pytorch_transformers/modeling_bert.py +++ b/pytorch_transformers/modeling_bert.py @@ -852,7 +852,7 @@ class BertForMaskedLM(BertPreTrainedModel): sequence_output = outputs[0] prediction_scores = self.cls(sequence_output) - outputs = (prediction_scores,) + outputs[2:] # Add hidden states and attention is they are here + outputs = (prediction_scores,) + outputs[2:] # Add hidden states and attention if they are here if masked_lm_labels is not None: loss_fct = CrossEntropyLoss(ignore_index=-1) masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), masked_lm_labels.view(-1))