Update attention.py

2023-08-08 17:47:08 +08:00 · 2023-08-08 17:47:08 +08:00 · b64bba71e1
parent 827a216155
commit b64bba71e1
1 changed files with 1 additions and 1 deletions
--- a/mmpretrain/models/utils/attention.py
+++ b/mmpretrain/models/utils/attention.py
@ -38,7 +38,7 @@ def scaled_dot_product_attention_pyimpl(query,
        attn_mask = torch.ones(
            query.size(-2), key.size(-2), dtype=torch.bool).tril(diagonal=0)
    if attn_mask is not None and attn_mask.dtype == torch.bool:
-        attn_mask = attn_mask.masked_fill(not attn_mask, -float('inf'))
+        attn_mask = attn_mask.masked_fill(~attn_mask, -float('inf'))

    attn_weight = query @ key.transpose(-2, -1) / scale
    if attn_mask is not None: