From bee9576f03f483a269e917da075d96c92faf35f1 Mon Sep 17 00:00:00 2001
From: Fredo Guan <fredo.guan@hotmail.com>
Date: Tue, 17 Jan 2023 11:11:19 -0800
Subject: [PATCH] Update metaformers.py

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

diff --git a/timm/models/metaformers.py b/timm/models/metaformers.py
index b7435128..d5bf9692 100644
--- a/timm/models/metaformers.py
+++ b/timm/models/metaformers.py
@@ -809,11 +809,12 @@ class MetaFormer(nn.Module):
         if pre_logits:
             return x
         
-        x = self.global_pool(x)
-        x = x.squeeze()
-        x = self.norm(x)
+        #x = self.global_pool(x)
+        #x = x.squeeze()
+        #x = self.norm(x)
         # (B, H, W, C) -> (B, C)
-        x = self.head(x)
+        #x = self.head(x)
+        x=self.head(self.norm(x.mean([1, 2])))
         return x
         
     def forward_features(self, x):