From 4d541f516fade15d39b35065155239fd3bf0299a Mon Sep 17 00:00:00 2001
From: George Mihaila <georgemihaila@my.unt.edu>
Date: Tue, 6 Oct 2020 08:12:04 -0500
Subject: [PATCH] fix return dicitonary labels from masked_lm_labels to labels
 (#7595)

---
 src/transformers/data/data_collator.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/transformers/data/data_collator.py b/src/transformers/data/data_collator.py
index 485045b609..4cf1b8919c 100644
--- a/src/transformers/data/data_collator.py
+++ b/src/transformers/data/data_collator.py
@@ -446,7 +446,7 @@ class DataCollatorForNextSentencePrediction:
             "input_ids": input_ids,
             "attention_mask": self._tensorize_batch(attention_masks),
             "token_type_ids": self._tensorize_batch(segment_ids),
-            "masked_lm_labels": mlm_labels if self.mlm else None,
+            "labels": mlm_labels if self.mlm else None,
             "next_sentence_label": torch.tensor(nsp_labels),
         }
         if self.mlm: