Update metaformers.py

2 years ago · 808f4a7ebd
parent bee9576f03
commit 808f4a7ebd
1 changed files with 1 additions and 1 deletions
--- a/timm/models/metaformers.py
+++ b/timm/models/metaformers.py
@ -814,7 +814,7 @@ class MetaFormer(nn.Module):
        #x = self.norm(x)
        # (B, H, W, C) -> (B, C)
        #x = self.head(x)
-        x=self.head(self.norm(x.mean([1, 2])))
+        x=self.head(self.norm(x.mean([2, 3])))
        return x
        
    def forward_features(self, x):