From 077478637d9e45c02a64f9cd1877bcc5b708b70d Mon Sep 17 00:00:00 2001 From: Sylvain Gugger <35901082+sgugger@users.noreply.github.com> Date: Mon, 26 Oct 2020 09:23:12 -0400 Subject: [PATCH] Fix label name in DataCollatorForNextSentencePrediction test (#8048) --- tests/test_data_collator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_data_collator.py b/tests/test_data_collator.py index f71e864752..ff45e87d6c 100644 --- a/tests/test_data_collator.py +++ b/tests/test_data_collator.py @@ -175,7 +175,7 @@ class DataCollatorIntegrationTest(unittest.TestCase): total_samples = batch["input_ids"].shape[0] self.assertEqual(batch["input_ids"].shape, torch.Size((total_samples, 512))) self.assertEqual(batch["token_type_ids"].shape, torch.Size((total_samples, 512))) - self.assertEqual(batch["masked_lm_labels"].shape, torch.Size((total_samples, 512))) + self.assertEqual(batch["labels"].shape, torch.Size((total_samples, 512))) self.assertEqual(batch["next_sentence_label"].shape, torch.Size((total_samples,))) @slow