"git@developer.sourcefind.cn:chenpangpang/ComfyUI.git" did not exist on "ebfcf0a9c91da34feec1d3ee6edd41a0302a637c"
Unverified Commit 4d541f51 authored by George Mihaila's avatar George Mihaila Committed by GitHub
Browse files

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

parent 8d2c248d
...@@ -446,7 +446,7 @@ class DataCollatorForNextSentencePrediction: ...@@ -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:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment