fix all SDPA dropouts

timm/models/beit.py

@@ -155,7 +155,7 @@ class Attention(nn.Module):
             x = F.scaled_dot_product_attention(
                 q, k, v,
                 attn_mask=rel_pos_bias,
-                dropout_p=self.attn_drop.p,
+                dropout_p=self.attn_drop.p if self.training else 0.,
             )
         else:
             q = q * self.scale

timm/models/cait.py

@@ -50,7 +50,7 @@ class ClassAttn(nn.Module):
         if self.fused_attn:
             x_cls = 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.,
             )
         else:
             q = q * self.scale

timm/models/eva.py

@@ -126,7 +126,7 @@ class EvaAttention(nn.Module):
             x = F.scaled_dot_product_attention(
                 q, k, v,
                 attn_mask=attn_mask,
-                dropout_p=self.attn_drop.p,
+                dropout_p=self.attn_drop.p if self.training else 0.,
             )
         else:
             q = q * self.scale

timm/models/fastvit.py

@@ -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 if self.training else 0.0,
+                dropout_p=self.attn_drop.p if self.training else 0.,
             )
         else:
             q = q * self.scale

timm/models/maxxvit.py

@@ -190,7 +190,7 @@ class Attention2d(nn.Module):
                 k.transpose(-1, -2).contiguous(),
                 v.transpose(-1, -2).contiguous(),
                 attn_mask=attn_bias,
-                dropout_p=self.attn_drop.p,
+                dropout_p=self.attn_drop.p if self.training else 0.,
             ).transpose(-1, -2).reshape(B, -1, H, W)
         else:
             q = q * self.scale
@@ -259,7 +259,7 @@ class AttentionCl(nn.Module):
             x = torch.nn.functional.scaled_dot_product_attention(
                 q, k, v,
                 attn_mask=attn_bias,
-                dropout_p=self.attn_drop.p,
+                dropout_p=self.attn_drop.p if self.training else 0.,
             )
         else:
             q = q * self.scale

timm/models/metaformer.py

@@ -198,7 +198,7 @@ class Attention(nn.Module):
         if self.fused_attn:
             x = F.scaled_dot_product_attention(
                 q, k, v,
-                dropout_p=self.attn_drop.p,
+                dropout_p=self.attn_drop.p if self.training else 0.,
             )
         else:
             attn = (q @ k.transpose(-2, -1)) * self.scale

timm/models/nest.py

@@ -66,7 +66,7 @@ class Attention(nn.Module):
         q, k, v = qkv.unbind(0)  # make torchscript happy (cannot use tensor as tuple)
 
         if self.fused_attn:
-            x = F.scaled_dot_product_attention(q, k, v, dropout_p=self.attn_drop.p)
+            x = F.scaled_dot_product_attention(q, k, v, dropout_p=self.attn_drop.p if self.training else 0.)
         else:
             q = q * self.scale
             attn = q @ k.transpose(-2, -1) # (B, H, T, N, N)

timm/models/pvt_v2.py

@@ -130,7 +130,7 @@ class Attention(nn.Module):
         k, v = kv.unbind(0)
 
         if self.fused_attn:
-            x = F.scaled_dot_product_attention(q, k, v, dropout_p=self.attn_drop.p)
+            x = F.scaled_dot_product_attention(q, k, v, dropout_p=self.attn_drop.p if self.training else 0.)
         else:
             q = q * self.scale
             attn = q @ k.transpose(-2, -1)

timm/models/swin_transformer.py

@@ -164,7 +164,7 @@ class WindowAttention(nn.Module):
             x = torch.nn.functional.scaled_dot_product_attention(
                 q, k, v,
                 attn_mask=attn_mask,
-                dropout_p=self.attn_drop.p,
+                dropout_p=self.attn_drop.p if self.training else 0.,
             )
         else:
             q = q * self.scale

timm/models/twins.py

@@ -75,7 +75,7 @@ class LocallyGroupedAttn(nn.Module):
         if self.fused_attn:
             x = F.scaled_dot_product_attention(
                 q, k, v,
-                dropout_p=self.attn_drop.p,
+                dropout_p=self.attn_drop.p if self.training else 0.,
             )
         else:
             q = q * self.scale
@@ -172,7 +172,7 @@ class GlobalSubSampleAttn(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.,
             )
         else:
             q = q * self.scale

timm/models/visformer.py

@@ -95,7 +95,7 @@ class Attention(nn.Module):
         if self.fused_attn:
             x = torch.nn.functional.scaled_dot_product_attention(
                 q.contiguous(), k.contiguous(), v.contiguous(),
-                dropout_p=self.attn_drop.p,
+                dropout_p=self.attn_drop.p if self.training else 0.,
             )
         else:
             attn = (q @ k.transpose(-2, -1)) * self.scale

timm/models/vision_transformer.py

@@ -85,7 +85,7 @@ class Attention(nn.Module):
         if self.fused_attn:
             x = F.scaled_dot_product_attention(
                 q, k, v,
-                dropout_p=self.attn_drop.p,
+                dropout_p=self.attn_drop.p if self.training else 0.,
             )
         else:
             q = q * self.scale
@@ -285,7 +285,7 @@ class ParallelScalingBlock(nn.Module):
         if self.fused_attn:
             x_attn = F.scaled_dot_product_attention(
                 q, k, v,
-                dropout_p=self.attn_drop.p,
+                dropout_p=self.attn_drop.p if self.training else 0.,
             )
         else:
             q = q * self.scale

timm/models/vision_transformer_relpos.py

@@ -71,7 +71,7 @@ class RelPosAttention(nn.Module):
             x = torch.nn.functional.scaled_dot_product_attention(
                 q, k, v,
                 attn_mask=attn_bias,
-                dropout_p=self.attn_drop.p,
+                dropout_p=self.attn_drop.p if self.training else 0.,
             )
         else:
             q = q * self.scale

timm/models/vision_transformer_sam.py

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