From 4ecf422cd3c693c21ff018ea88868642b4e6cf90 Mon Sep 17 00:00:00 2001 From: Fredo Guan Date: Sat, 10 Dec 2022 05:46:27 -0800 Subject: [PATCH] Update davit.py --- timm/models/davit.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/timm/models/davit.py b/timm/models/davit.py index 91e5f50a..18cf5af9 100644 --- a/timm/models/davit.py +++ b/timm/models/davit.py @@ -576,8 +576,8 @@ class DaViT(nn.Module): def forward_features(self, x): x, sizes = self.forward_network(x) # take final feature and norm - x = self.norms(x) - H, W = size + x = self.norms(x[-1]) + H, W = sizes[-1] x = x.view(-1, H, W, self.embed_dims[-1]).permute(0, 3, 1, 2).contiguous() return x