Add first_conv for mambaout

huggingface · Sep 14, 2024 · 36ecd37 · 36ecd37
1 parent 02d29de
commit 36ecd37
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/timm/models/mambaout.py b/timm/models/mambaout.py
@@ -314,7 +314,6 @@ def __init__(
             drop_path_rate=0.,
             drop_rate=0.,
             head_fn='default',
-            **kwargs,
     ):
         super().__init__()
         self.num_classes = num_classes
@@ -461,7 +460,8 @@ def _cfg(url='', **kwargs):
         'url': url,
         'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': (7, 7),
         'crop_pct': 1.0, 'interpolation': 'bicubic',
-        'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD, 'classifier': 'head.fc',
+        'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD,
+        'first_conv': 'stem.conv1', 'classifier': 'head.fc',
         **kwargs
     }