Update tests/test_model.py

mosaicml · Sep 16, 2023 · 2846d29 · 2846d29
1 parent 3c2a975
commit 2846d29
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_model.py b/tests/test_model.py
@@ -1341,7 +1341,7 @@ def test_forward_with_output_attentions_and_output_hidden_states(
 
         if output_attentions:
             assert len(outputs.attentions) == n_layers
-            assert all(attn.shape == (1, 4, 2048, 2048) for attn in outputs.attentions)
+            assert all(attn.shape == (1, 4, 3, 3) for attn in outputs.attentions)
         if output_hidden_states:
             assert len(outputs.hidden_states) == n_layers + 1