From 39274bd45e78b8ead0509367f800121f0d7c25f4 Mon Sep 17 00:00:00 2001
From: Fredo Guan <fredo.guan@hotmail.com>
Date: Sun, 5 Feb 2023 02:06:58 -0800
Subject: [PATCH] Update metaformers.py

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

diff --git a/timm/models/metaformers.py b/timm/models/metaformers.py
index 1af3770f..db9421dc 100644
--- a/timm/models/metaformers.py
+++ b/timm/models/metaformers.py
@@ -582,7 +582,7 @@ class MetaFormer(nn.Module):
     
     def forward_head(self, x, pre_logits: bool = False):
         # NOTE nn.Sequential in head broken down since can't call head[:-1](x) in torchscript :(
-        x = self.head.global_pool(x.permute(0, 3, 1, 2))
+        x = self.head.global_pool(x)
         x = self.head.flatten(x)
         x = self.head.norm(x)
         return x if pre_logits else self.head.fc(x)
@@ -593,7 +593,7 @@ class MetaFormer(nn.Module):
             x = checkpoint_seq(self.stages, x)
         else:
             x = self.stages(x)
-        x = self.norm_pre(x)
+        x = self.norm_pre(x).permute(0, 3, 1, 2)
         return x 
 
     def forward(self, x):