add main test file

jshuadvd · Jun 10, 2024 · d14436d · d14436d
1 parent d90c718
commit d14436d
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/tests/test_main.py b/tests/test_main.py
@@ -9,3 +9,10 @@
     progressive_extension,
     short_context_recovery,
 )
+
+
+def test_load_data():
+    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+    data = load_data("data/raw/enwik8.gz", tokenizer, 65536)
+    assert len(data) > 0
+    assert isinstance(data[0], torch.Tensor)