mirror of
https://github.com/huggingface/pytorch-image-models.git
synced 2025-06-03 15:01:08 +08:00
Update metaformers.py
This commit is contained in:
parent
4cfecf8acb
commit
1145720bec
@ -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])
|
||||
|
Loading…
x
Reference in New Issue
Block a user