Add assertion message

Co-authored-by: Madeesh Kannan <[email protected]>
explosion · Feb 8, 2024 · eefe900 · eefe900
1 parent 130df32
commit eefe900
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/curated_transformers/layers/attention.py b/curated_transformers/layers/attention.py
@@ -764,7 +764,7 @@ def forward(
             #
             # Doing this properly requires a redesign of our AttentionMask
             # class.
-            assert attention_mask.bool_mask.size(-2) == 1
+            assert attention_mask.bool_mask.size(-2) == 1, "Torch SDP does not support attention masks with non-broadcastable query length yet"
             return torch.where(
                 attention_mask.bool_mask.transpose(-1, -2), attn_values, 0.0
             )