From 4469b4fbee8389bbd9f91d3f3ffee5bd1731c3f7 Mon Sep 17 00:00:00 2001 From: Fredo Guan Date: Thu, 8 Dec 2022 07:22:41 -0800 Subject: [PATCH] Update davit.py --- timm/models/davit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/timm/models/davit.py b/timm/models/davit.py index 539086c4..e3d9d719 100644 --- a/timm/models/davit.py +++ b/timm/models/davit.py @@ -439,7 +439,7 @@ class DaViT(nn.Module): ]) self.stages.add_module(f'stage_{stage_id}', stage) - self.feature_info += [dict(num_chs=self.embed_dims[item], reduction=2, module=f'stages.stage_{stage_id}')] + self.feature_info += [dict(num_chs=self.embed_dims[stage_id], reduction=2, module=f'stages.stage_{stage_id}')] self.norms = norm_layer(self.num_features) self.head = ClassifierHead(self.num_features, num_classes, pool_type=global_pool, drop_rate=drop_rate)