mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

eva.py: fixed bug in applying attention mask

Browse Source 
The mask should be applied before the softmax.
This commit is contained in:
Feraidoon Mehri 2024-07-19 15:12:04 +03:30 committed by GitHub
parent 7160af4a24
commit 4cca568bd8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
timm/models/eva.py
View File

@ -134,10 +134,12 @@ class EvaAttention(nn.Module):
else: else:
q = q * self.scale q = q * self.scale
attn = (q @ k.transpose(-2, -1)) attn = (q @ k.transpose(-2, -1))
attn = attn.softmax(dim=-1)
if attn_mask is not None: if attn_mask is not None:
attn_mask = attn_mask.to(torch.bool) attn_mask = attn_mask.to(torch.bool)
attn = attn.masked_fill(~attn_mask[:, None, None, :], float("-inf")) attn = attn.masked_fill(~attn_mask[:, None, None, :], float("-inf"))
attn = attn.softmax(dim=-1)
attn = self.attn_drop(attn) attn = self.attn_drop(attn)
x = attn @ v x = attn @ v