diff --git a/timm/models/davit.py b/timm/models/davit.py index 33fd403e..30d30879 100644 --- a/timm/models/davit.py +++ b/timm/models/davit.py @@ -443,7 +443,7 @@ class DaViT(nn.Module): self.feature_info += [dict(num_chs=self.embed_dims[stage_id], reduction = 2, module=f'stage_{stage_id}')] - self.norm = norm_layer(self.num_features) + self.norms = norm_layer(self.num_features) self.head = ClassifierHead(self.num_features, num_classes, pool_type=global_pool, drop_rate=drop_rate) self.apply(self._init_weights) @@ -507,7 +507,7 @@ class DaViT(nn.Module): def forward_features(self, x): x, sizes = self.forward_network(x) # take final feature and norm - x = self.norm(x[-1]) + 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