Fix visformer in_chans stem handling

Author: rwightman

tests/test_models.py

@@ -190,7 +190,7 @@ def test_model_features_pretrained(model_name, batch_size):
 def test_model_forward_torchscript(model_name, batch_size):
     """Run a single forward pass with each model"""
     input_size = _get_input_size(model_name=model_name, target=TARGET_JIT_SIZE)
-    if max(input_size) > MAX_JIT_SIZE:  # NOTE using MAX_FWD_SIZE as the final limit is intentional
+    if max(input_size) > MAX_JIT_SIZE:
         pytest.skip("Fixed input size model > limit.")
 
     with set_scriptable(True):

timm/models/visformer.py

@@ -26,7 +26,7 @@ def _cfg(url='', **kwargs):
         'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None,
         'crop_pct': .9, 'interpolation': 'bicubic', 'fixed_input_size': True,
         'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD,
-        'first_conv': 'patch_embed.proj', 'classifier': 'head',
+        'first_conv': 'stem.0', 'classifier': 'head',
         **kwargs
     }
 
@@ -183,7 +183,7 @@ def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, in
                 img_size //= 8
             else:
                 self.stem = nn.Sequential(
-                    nn.Conv2d(3, self.init_channels, 7, stride=2, padding=3, bias=False),
+                    nn.Conv2d(in_chans, self.init_channels, 7, stride=2, padding=3, bias=False),
                     nn.BatchNorm2d(self.init_channels),
                     nn.ReLU(inplace=True)
                 )