mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Enable features_only properly

Browse Source
This commit is contained in:
Ross Wightman 2024-06-04 16:57:16 -07:00
parent 1b66ec7cf3
commit 58591a97f7
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
timm/models/vitamin.py
View File

@ -279,6 +279,7 @@ class GeGluMlp(nn.Module):
def _create_vitamin(variant, pretrained=False, embed_cfg=None, **kwargs): def _create_vitamin(variant, pretrained=False, embed_cfg=None, **kwargs):
out_indices = kwargs.pop('out_indices', 3)
assert embed_cfg is not None assert embed_cfg is not None
backbone = MbConvStages(cfg=embed_cfg) backbone = MbConvStages(cfg=embed_cfg)
kwargs['embed_layer'] = partial(HybridEmbed, backbone=backbone, proj=False) kwargs['embed_layer'] = partial(HybridEmbed, backbone=backbone, proj=False)
@ -289,6 +290,7 @@ def _create_vitamin(variant, pretrained=False, embed_cfg=None, **kwargs):
variant, variant,
pretrained, pretrained,
pretrained_filter_fn=checkpoint_filter_fn, pretrained_filter_fn=checkpoint_filter_fn,
feature_cfg=dict(out_indices=out_indices, feature_cls='getter'),
**kwargs, **kwargs,
) )