From c2523be634a94da2b1b2a43c11b25827a0de990d Mon Sep 17 00:00:00 2001
From: Ding Yiwei <16083536+dingyiwei@users.noreply.github.com>
Date: Mon, 15 Nov 2021 17:06:18 +0800
Subject: [PATCH] Replace 2 `transpose()` with 1 `permute` in
 TransformerBlock()` (#5645)

---
 models/common.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/models/common.py b/models/common.py
index 3ea7ba5..3930c8e 100644
--- a/models/common.py
+++ b/models/common.py
@@ -86,8 +86,8 @@ class TransformerBlock(nn.Module):
         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):