Update tests/test_model.py

Co-authored-by: Daniel King <[email protected]>
mosaicml · Sep 15, 2023 · 2fad7a5 · 2fad7a5
1 parent 950ef42
commit 2fad7a5
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 outputs.attentions[0] is not None
+            assert all(attn.shape == (1, 4, 2048, 2048) for attn in outputs.attentions)
         if output_hidden_states:
             assert len(outputs.hidden_states) == n_layers + 1