From 116345b4d39218a379d2928c44625fd287011a18 Mon Sep 17 00:00:00 2001
From: Fredo Guan <fredo.guan@hotmail.com>
Date: Sat, 7 Jan 2023 21:31:01 -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 257d4e74..a7404cf9 100644
--- a/timm/models/metaformers.py
+++ b/timm/models/metaformers.py
@@ -627,10 +627,11 @@ class MetaFormer(nn.Module):
         return {'norm'}
 
     def forward_features(self, x):
-        for i in range(self.num_stage):
-            x = self.downsample_layers[i](x)
-            x = self.stages[i](x)
-        return self.norm(x.mean([1, 2])) # (B, H, W, C) -> (B, C)
+        x = self.stages(x)
+        x = x.mean([1,2]) # TODO use adaptive pool instead of mean
+        x = self.norm(x)
+        # (B, H, W, C) -> (B, C)
+        return x 
 
     def forward(self, x):
         x = self.forward_features(x)