From 58591a97f7607f9b860bd7a77cc69f29a008057d Mon Sep 17 00:00:00 2001
From: Ross Wightman <rwightman@gmail.com>
Date: Tue, 4 Jun 2024 16:57:16 -0700
Subject: [PATCH] Enable features_only properly

---
 timm/models/vitamin.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/timm/models/vitamin.py b/timm/models/vitamin.py
index 71d3b674..7c2c8735 100644
--- a/timm/models/vitamin.py
+++ b/timm/models/vitamin.py
@@ -279,6 +279,7 @@ class GeGluMlp(nn.Module):
 
 
 def _create_vitamin(variant, pretrained=False, embed_cfg=None, **kwargs):
+    out_indices = kwargs.pop('out_indices', 3)
     assert embed_cfg is not None
     backbone = MbConvStages(cfg=embed_cfg)
     kwargs['embed_layer'] = partial(HybridEmbed, backbone=backbone, proj=False)
@@ -289,6 +290,7 @@ def _create_vitamin(variant, pretrained=False, embed_cfg=None, **kwargs):
         variant,
         pretrained,
         pretrained_filter_fn=checkpoint_filter_fn,
+        feature_cfg=dict(out_indices=out_indices, feature_cls='getter'),
         **kwargs,
     )