Update metaformers.py

2 years ago · 1145720bec
parent 4cfecf8acb
commit 1145720bec
1 changed files with 1 additions and 1 deletions
--- 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])