Update passkey_retrivial.py

dvlab-research · Nov 15, 2023 · 0e0622b · 0e0622b
1 parent 26a6b69
commit 0e0622b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/passkey_retrivial.py b/passkey_retrivial.py
@@ -93,7 +93,7 @@ def main(args):
     print("base model", args.base_model)
 
     if args.flash_attn:
-        replace_llama_attn()
+        replace_llama_attn(use_full=True)
 
     # Set RoPE scaling factor
     config = transformers.AutoConfig.from_pretrained(