From 61e1cbc100e149774fcc318227e932fea297535d Mon Sep 17 00:00:00 2001 From: dingyiwei Date: Mon, 15 Nov 2021 03:04:24 +0000 Subject: [PATCH] Use permute() instead of 2 transpose() --- models/common.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/common.py b/models/common.py index 3ea7ba5477a6..3930c8e7b2df 100644 --- 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):