Additional test fixes

stanfordnlp · Jan 29, 2025 · 5294462 · 5294462
1 parent 3f30793
commit 5294462
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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,
     )