fix mask padding trim missing from non-memory-efficient attention path

src/audiocraft/modules/transformer.py

@@ -399,15 +399,18 @@ class StreamingMultiheadAttention(StreamingModule):
                     v = expand_repeated_kv(v, self.kv_repeat, self.memory_efficient)
             if self.attention_as_float32:
                 q, k, v = [x.float() for x in [q, k, v]]
-            if self.memory_efficient:
-                if custom_attn_mask:
-                    # When using a custom attn mask:
-                    # Move to query's device, repeat for each sample, remove align8 padding
-                    seq_len = query.shape[1]
-                    attn_mask = attn_mask.to(q.dtype)
-                    attn_mask = attn_mask.repeat((q.shape[0], 1, 1, 1))
-                    attn_mask = attn_mask[..., :seq_len, :seq_len]
+            # Original: this block was only reachable inside `if self.memory_efficient:`.
+            # Moved out so the plain (non-memory-efficient) path below also gets the
+            # align8-padding trim -- it needs the same seq_len match as the mask.
+            if custom_attn_mask:
+                # When using a custom attn mask:
+                # Move to query's device, repeat for each sample, remove align8 padding
+                seq_len = query.shape[1]
+                attn_mask = attn_mask.to(q.dtype)
+                attn_mask = attn_mask.repeat((q.shape[0], 1, 1, 1))
+                attn_mask = attn_mask[..., :seq_len, :seq_len]
 
+            if self.memory_efficient:
                 p = self.dropout if self.training else 0
                 if _efficient_attention_backend == 'torch':
                     x = torch.nn.functional.scaled_dot_product_attention(