Update multihead_attention.py

2023-04-23 18:45:48 -07:00 · 2023-04-23 18:45:48 -07:00 · 62cedb9c8f
commit 62cedb9c8f
parent d4a62ccfb5
1 changed files with 3 additions and 0 deletions
--- a/torchscale/component/multihead_attention.py
+++ b/torchscale/component/multihead_attention.py
@ -133,6 +133,9 @@ class MultiheadAttention(nn.Module):
            attn = F.scaled_dot_product_attention(
                q, k, v, attn_mask, self.dropout_module.p
            )
+            # attn: B,H,T,E (Batch, Heads, Tgt_Len, Dim)
+            # Permute to B,T,H,E, and then flatten to B,T,D
+            attn = attn.permute(0, 2, 1, 3).flatten(2)
            attn_weights = None
        else:
            q *= self.scaling