fastvit: don't dropout in eval

YassineYousfi · rwightman · commit b500cae4c537 · 2023-10-05T08:58:41.000-07:00
diff --git a/timm/models/fastvit.py b/timm/models/fastvit.py
@@ -514,7 +514,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         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

Original file line number	Diff line number	Diff line change
`@@ -514,7 +514,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:`
`514`	`514`	`if self.fused_attn:`
`515`	`515`	`x = torch.nn.functional.scaled_dot_product_attention(`
`516`	`516`	`q, k, v,`
`517`		`- dropout_p=self.attn_drop.p,`
	`517`	`+ dropout_p=self.attn_drop.p if self.training else 0.0,`
`518`	`518`	`)`
`519`	`519`	`else:`
`520`	`520`	`q = q * self.scale`