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

fastvit: don't dropout in eval

Browse Source
This commit is contained in:
Yassine 2023-10-03 11:57:24 -07:00 committed by Ross Wightman
parent 054c763fca
commit b500cae4c5
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
timm/models/fastvit.py
View File

@ -514,7 +514,7 @@ class Attention(nn.Module):
if self.fused_attn:
x = torch.nn.functional.scaled_dot_product_attention(
q, k, v,
dropout_p=self.attn_drop.p,
dropout_p=self.attn_drop.p if self.training else 0.0,
)
else:
q = q * self.scale