From 0bde1c12181ee0828f2a405b6dd58b73f5edd9c2 Mon Sep 17 00:00:00 2001 From: Fredo Guan Date: Sat, 7 Jan 2023 21:43:56 -0800 Subject: [PATCH] Update metaformers.py --- timm/models/metaformers.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/timm/models/metaformers.py b/timm/models/metaformers.py index 656f6294..99b9b53f 100644 --- a/timm/models/metaformers.py +++ b/timm/models/metaformers.py @@ -601,10 +601,11 @@ class MetaFormer(nn.Module): res_scale_init_value=res_scale_init_values[i], ) for j in range(depths[i])] ) + stages.append(downsample_layers[i]) stages.append(stage) cur += depths[i] - self.stages = nn.Sequential(zip(*downsample_layers, *stages)) + self.stages = nn.Sequential(*stages) self.norm = output_norm(dims[-1])