Update metaformers.py

This commit is contained in:
Fredo Guan 2023-01-17 11:13:14 -08:00
parent bee9576f03
commit 808f4a7ebd

View File

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