SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

fix return dicitonary labels from masked_lm_labels to labels (#7595)

Browse Source
This commit is contained in:
George Mihaila
2020-10-06 08:12:04 -05:00
committed by GitHub
parent 8d2c248df7
commit 4d541f516f
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/transformers/data/data_collator.py
View File

@@ -446,7 +446,7 @@ class DataCollatorForNextSentencePrediction:
"input_ids": input_ids, "input_ids": input_ids,
"attention_mask": self._tensorize_batch(attention_masks), "attention_mask": self._tensorize_batch(attention_masks),
"token_type_ids": self._tensorize_batch(segment_ids), "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), "next_sentence_label": torch.tensor(nsp_labels),
} }
if self.mlm: if self.mlm: