From b81e79aae9579a6868e139780ad064d73693c2d3 Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Tue, 28 Sep 2021 16:38:41 -0700 Subject: [PATCH] Fix bottleneck attn transpose typo, hopefully these train better now.. --- timm/models/layers/bottleneck_attn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/timm/models/layers/bottleneck_attn.py b/timm/models/layers/bottleneck_attn.py index c0c619cc..bf6af675 100644 --- a/timm/models/layers/bottleneck_attn.py +++ b/timm/models/layers/bottleneck_attn.py @@ -122,7 +122,7 @@ class BottleneckAttn(nn.Module): attn_logits = attn_logits + self.pos_embed(q) # B, num_heads, H * W, H * W attn_out = attn_logits.softmax(dim=-1) - attn_out = (attn_out @ v).transpose(1, 2).reshape(B, self.dim_out, H, W) # B, dim_out, H, W + attn_out = (attn_out @ v).transpose(-1, -2).reshape(B, self.dim_out, H, W) # B, dim_out, H, W attn_out = self.pool(attn_out) return attn_out