mirror of
https://github.com/huggingface/pytorch-image-models.git
synced 2025-06-03 15:01:08 +08:00
fastvit: don't dropout in eval
This commit is contained in:
parent
054c763fca
commit
b500cae4c5
@ -514,7 +514,7 @@ class Attention(nn.Module):
|
|||||||
if self.fused_attn:
|
if self.fused_attn:
|
||||||
x = torch.nn.functional.scaled_dot_product_attention(
|
x = torch.nn.functional.scaled_dot_product_attention(
|
||||||
q, k, v,
|
q, k, v,
|
||||||
dropout_p=self.attn_drop.p,
|
dropout_p=self.attn_drop.p if self.training else 0.0,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
q = q * self.scale
|
q = q * self.scale
|
||||||
|
Loading…
x
Reference in New Issue
Block a user