ViTamin in_chans !=3 weight load fix

huggingface · Jun 8, 2024 · 7702d9a · 7702d9a
1 parent 5517b05
commit 7702d9a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/timm/models/vitamin.py b/timm/models/vitamin.py
@@ -281,7 +281,7 @@ def forward(self, x):
 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)
+    backbone = MbConvStages(cfg=embed_cfg, in_chans=kwargs.get('in_chans', 3))
     kwargs['embed_layer'] = partial(HybridEmbed, backbone=backbone, proj=False)
     kwargs.setdefault('patch_size', 1)  # default patch size for hybrid models if not set