From 1145720bec9b0facc3b26b1421c288a9f71eb7ae Mon Sep 17 00:00:00 2001 From: Fredo Guan Date: Sun, 8 Jan 2023 11:21:58 -0800 Subject: [PATCH] Update metaformers.py --- timm/models/metaformers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/timm/models/metaformers.py b/timm/models/metaformers.py index 2ba67b0a..0e9a58cb 100644 --- a/timm/models/metaformers.py +++ b/timm/models/metaformers.py @@ -668,7 +668,7 @@ class MetaFormer(nn.Module): ) stages.append(stage) cur += depths[i] - self.feature_info += [dict(num_chs=dims[stage_id], reduction=2, module=f'stages.{stage_id}')] + self.feature_info += [dict(num_chs=dims[i], reduction=2, module=f'stages.{i}')] self.stages = nn.Sequential(*stages) self.norm = output_norm(dims[-1])