@@ -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): |