mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Update metaformers.py

Browse Source
This commit is contained in:
Fredo Guan 2023-01-08 12:13:14 -08:00
parent ddfdb543dc
commit 8568bc7b6a
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
timm/models/metaformers.py
View File

@ -236,9 +236,13 @@ class Downsampling(nn.Module):
self.post_norm = post_norm(out_channels) if post_norm else nn.Identity()
def forward(self, x):
print(x.shape)
x = self.pre_norm(x)
print(x.shape)
x = self.conv(x)
print(x.shape)
x = self.post_norm(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
print(x.shape)
return x
class Scale(nn.Module):