Replace 2 transpose() with 1 permute in TransformerBlock()` (#5645)

ultralytics · Nov 15, 2021 · c2523be · c2523be
1 parent 540ef0d
commit c2523be
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/models/common.py b/models/common.py
@@ -86,8 +86,8 @@ def forward(self, x):
         if self.conv is not None:
             x = self.conv(x)
         b, _, w, h = x.shape
-        p = x.flatten(2).unsqueeze(0).transpose(0, 3).squeeze(3)
-        return self.tr(p + self.linear(p)).unsqueeze(3).transpose(0, 3).reshape(b, self.c2, w, h)
+        p = x.flatten(2).permute(2, 0, 1)
+        return self.tr(p + self.linear(p)).permute(1, 2, 0).reshape(b, self.c2, w, h)
 
 
 class Bottleneck(nn.Module):