pretrain: fix with sample_packing=false (#1841)

2024-08-21 10:36:51 -07:00
parent 5aac4bc284
commit 649c19aba3
2 changed files with 3 additions and 3 deletions
--- a/tests/test_data.py
+++ b/tests/test_data.py
@@ -35,7 +35,7 @@ class TestEncodePretraining(unittest.TestCase):
                "hello, hello",
            ]
        }
-        result = encode_pretraining(self.tokenizer, self.max_tokens, examples["text"])
+        result = encode_pretraining(self.tokenizer, self.max_tokens, examples)

        self.assertEqual(len(result["input_ids"]), 3)