From 529446283247d6f5911ffec804adac97f0b551ea Mon Sep 17 00:00:00 2001 From: PinetreePantry Date: Wed, 29 Jan 2025 14:29:19 -0800 Subject: [PATCH] Additional test fixes --- tests/utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/utils.py b/tests/utils.py index 62927c85..94dce25e 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -134,11 +134,12 @@ def GPT2_SELF_ATTENTION_RUN( ) CACHE_ACTIVATIONS[f"{i}.head_value_output"] = head_value - head_attention_value_output, attn_weights = eager_attention_forward( + head_attention_value_output = eager_attention_forward( module=self_attn, query=head_query, key=head_key, value=head_value, + attention_mask=None, )