Add num_beams to make tests pass

mlfoundations · Dec 15, 2023 · b34e6ac · b34e6ac
1 parent b5f9beb
commit b34e6ac
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/tests/test_generate_kv_cache_time.py b/tests/test_generate_kv_cache_time.py
@@ -33,6 +33,7 @@ def test_generate_kv_cache(wiki_page, context_len, max_gen_len):
             "qk_norm": False,
             "positional_embedding_type": "rotary",
             "ffn_type": "swiglu",
+            "num_beams": 1,
         }
     )
 

diff --git a/tests/test_generate_load_kv_cache_equal.py b/tests/test_generate_load_kv_cache_equal.py
@@ -42,6 +42,7 @@ def args():
             "qk_norm": False,
             "positional_embedding_type": "head_rotary",
             "ffn_type": "swiglu",
+            "num_beams": 1,
         }
     )
     return args