Update multihead_attention.py

2023-04-23 18:28:08 -07:00 · 2023-04-23 18:28:08 -07:00 · d4a62ccfb5
commit d4a62ccfb5
parent 412a1a3878
1 changed files with 2 additions and 0 deletions
--- a/torchscale/component/multihead_attention.py
+++ b/torchscale/component/multihead_attention.py
@ -118,6 +118,8 @@ class MultiheadAttention(nn.Module):

        if attn_mask is not None:
            attn_mask = attn_mask.unsqueeze(0)
+        else:
+            attn_mask = torch.zeros(1, tgt_len, src_len, dtype=torch.float32, device=k.device)

        if key_padding_mask is not None:
            # Achieve same result with an additive mask