add flash_attention_causal_mask to run_lm_eval.py (#142)

HabanaAI · Apr 22, 2024 · 3601bb0 · 3601bb0
1 parent 0b2e152
commit 3601bb0
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/examples/text-generation/run_lm_eval.py b/examples/text-generation/run_lm_eval.py
@@ -87,6 +87,7 @@ def __init__(self, tokenizer, model, args, options):
                     "attn_softmax_bf16": self.options.attn_softmax_bf16,
                     "use_flash_attention": self.options.use_flash_attention,
                     "flash_attention_recompute": self.options.flash_attention_recompute,
+                    "flash_attention_causal_mask": self.options.flash_attention_causal_mask,
                 }
             )
         if args.warmup: