Temp fix for mask shape

2023-12-08 17:23:51 +01:00
parent a634b1ea09
commit 80d5d9299b
1 changed files with 1 additions and 1 deletions
@@ -73,7 +73,7 @@ class MultiHeadCrossAttention(nn.Module):
            q, k, v = map(lambda t: t.permute(0, 2, 1, 3),(q, k, v),)
            
            attn_mask = None
-            if mask is not None:
+            if mask is not None and len(mask) > 1:
                # This is probably wrong
                attn_mask = torch.zeros(
                    [1, q.shape[1], q.shape[2], v.shape[2]],