fix: [N, 1] pattern should take N to create 1d causal mask

facebookresearch · Nov 14, 2024 · 7814c5c · 7814c5c
1 parent 210e32a
commit 7814c5c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/xformers/components/attention/global_tokens.py b/xformers/components/attention/global_tokens.py
@@ -70,7 +70,7 @@ def __init__(
         self.force_sparsity = force_sparsity
 
         if causal:
-            self.attention_mask &= causal_1d_pattern(attention_query_mask.shape[1])
+            self.attention_mask &= causal_1d_pattern(attention_query_mask.shape[0])
 
         self.attention_mask = (
             sparsify(self.attention_mask)