Fix the short_context_recovery function call

jshuadvd · Jun 13, 2024 · ef76867 · ef76867
1 parent 77b87cd
commit ef76867
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/tests/test_main.py b/tests/test_main.py
@@ -155,7 +155,12 @@ def test_short_context_recovery():
         d_model=512, n_heads=8, num_layers=6, vocab_size=50257, max_len=65536
     )
     tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-    short_context_recovery(model, tokenizer)
+    base_length = 65536
+    lambda_factors_base = torch.ones(256)
+    n_hat_base = 50
+    short_context_recovery(
+        model, tokenizer, base_length, lambda_factors_base, n_hat_base
+    )
 
 
 # Testing the longrope_model_forward_with_extended_context function