From e4e43190ce3009af5205c75c1a9014d45cdcbefb Mon Sep 17 00:00:00 2001
From: Ross Wightman <rwightman@gmail.com>
Date: Mon, 8 May 2023 08:52:38 -0700
Subject: [PATCH 1/7] Add typing to all model entrypoint fns, add old cache
 check env var to builder

---
 timm/models/_builder.py                  |  10 +-
 timm/models/beit.py                      |  14 +-
 timm/models/byoanet.py                   |  30 +-
 timm/models/byobnet.py                   |  66 ++--
 timm/models/cait.py                      |  20 +-
 timm/models/coat.py                      |  16 +-
 timm/models/convit.py                    |  12 +-
 timm/models/convmixer.py                 |   6 +-
 timm/models/convnext.py                  |  50 +--
 timm/models/crossvit.py                  |  28 +-
 timm/models/cspnet.py                    | 368 +++++++++++------------
 timm/models/davit.py                     |  22 +-
 timm/models/deit.py                      |  32 +-
 timm/models/densenet.py                  |  12 +-
 timm/models/dla.py                       |  22 +-
 timm/models/dpn.py                       |  14 +-
 timm/models/edgenext.py                  |  10 +-
 timm/models/efficientformer.py           |   6 +-
 timm/models/efficientformer_v2.py        |   8 +-
 timm/models/efficientnet.py              | 206 ++++++-------
 timm/models/eva.py                       |  32 +-
 timm/models/focalnet.py                  |  24 +-
 timm/models/gcvit.py                     |  10 +-
 timm/models/ghostnet.py                  |   6 +-
 timm/models/hardcorenas.py               |  12 +-
 timm/models/hrnet.py                     |  22 +-
 timm/models/inception_resnet_v2.py       |   2 +-
 timm/models/inception_v3.py              |   2 +-
 timm/models/inception_v4.py              |   2 +-
 timm/models/levit.py                     |  40 +--
 timm/models/maxxvit.py                   | 112 +++----
 timm/models/metaformer.py                |  36 +--
 timm/models/mlp_mixer.py                 |  30 +-
 timm/models/mobilenetv3.py               |  40 +--
 timm/models/mobilevit.py                 |  20 +-
 timm/models/mvitv2.py                    |  16 +-
 timm/models/nasnet.py                    |   2 +-
 timm/models/nest.py                      |  12 +-
 timm/models/nfnet.py                     |  70 ++---
 timm/models/pit.py                       |  16 +-
 timm/models/pnasnet.py                   |   2 +-
 timm/models/pvt_v2.py                    |  14 +-
 timm/models/regnet.py                    |  72 ++---
 timm/models/res2net.py                   |  18 +-
 timm/models/resnest.py                   |  16 +-
 timm/models/resnet.py                    | 161 +++++-----
 timm/models/resnetv2.py                  |  32 +-
 timm/models/rexnet.py                    |  28 +-
 timm/models/selecsls.py                  |  72 ++---
 timm/models/senet.py                     |  18 +-
 timm/models/sequencer.py                 |   6 +-
 timm/models/sknet.py                     |  10 +-
 timm/models/swin_transformer.py          |  18 +-
 timm/models/swin_transformer_v2.py       |  24 +-
 timm/models/swin_transformer_v2_cr.py    |  32 +-
 timm/models/tnt.py                       |   4 +-
 timm/models/tresnet.py                   |   8 +-
 timm/models/twins.py                     |  12 +-
 timm/models/visformer.py                 |   4 +-
 timm/models/vision_transformer.py        | 122 ++++----
 timm/models/vision_transformer_hybrid.py |  28 +-
 timm/models/vision_transformer_relpos.py |  26 +-
 timm/models/volo.py                      |  22 +-
 timm/models/vovnet.py                    |  20 +-
 timm/models/xception.py                  |   3 +-
 timm/models/xception_aligned.py          |  10 +-
 timm/models/xcit.py                      |  62 ++--
 67 files changed, 1150 insertions(+), 1152 deletions(-)

diff --git a/timm/models/_builder.py b/timm/models/_builder.py
index f6a3f54182..cbf057b29b 100644
--- a/timm/models/_builder.py
+++ b/timm/models/_builder.py
@@ -22,7 +22,7 @@
 # Use set_pretrained_download_progress / set_pretrained_check_hash functions to toggle.
 _DOWNLOAD_PROGRESS = False
 _CHECK_HASH = False
-
+_USE_OLD_CACHE = int(os.environ.get('TIMM_USE_OLD_CACHE', 0)) > 0
 
 __all__ = ['set_pretrained_download_progress', 'set_pretrained_check_hash', 'load_custom_pretrained', 'load_pretrained',
            'pretrained_cfg_for_features', 'resolve_pretrained_cfg', 'build_model_with_cfg']
@@ -49,9 +49,11 @@ def _resolve_pretrained_source(pretrained_cfg):
             load_from = 'file'
             pretrained_loc = pretrained_file
         else:
-            # next, HF hub is prioritized unless a valid cached version of weights exists already
-            cached_url_valid = check_cached_file(pretrained_url) if pretrained_url else False
-            if hf_hub_id and has_hf_hub(necessary=True) and not cached_url_valid:
+            old_cache_valid = False
+            if _USE_OLD_CACHE:
+                # prioritized old cached weights if exists and env var enabled
+                old_cache_valid = check_cached_file(pretrained_url) if pretrained_url else False
+            if not old_cache_valid and hf_hub_id and has_hf_hub(necessary=True):
                 # hf-hub available as alternate weight source in default_cfg
                 load_from = 'hf-hub'
                 pretrained_loc = hf_hub_id
diff --git a/timm/models/beit.py b/timm/models/beit.py
index d3d77ba888..d4fdfb8df9 100644
--- a/timm/models/beit.py
+++ b/timm/models/beit.py
@@ -515,7 +515,7 @@ def _create_beit(variant, pretrained=False, **kwargs):
 
 
 @register_model
-def beit_base_patch16_224(pretrained=False, **kwargs):
+def beit_base_patch16_224(pretrained=False, **kwargs) -> Beit:
     model_args = dict(
         patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4,
         use_abs_pos_emb=False, use_rel_pos_bias=True, init_values=0.1)
@@ -524,7 +524,7 @@ def beit_base_patch16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def beit_base_patch16_384(pretrained=False, **kwargs):
+def beit_base_patch16_384(pretrained=False, **kwargs) -> Beit:
     model_args = dict(
         img_size=384, patch_size=16, embed_dim=768, depth=12, num_heads=12,
         use_abs_pos_emb=False, use_rel_pos_bias=True, init_values=0.1)
@@ -533,7 +533,7 @@ def beit_base_patch16_384(pretrained=False, **kwargs):
 
 
 @register_model
-def beit_large_patch16_224(pretrained=False, **kwargs):
+def beit_large_patch16_224(pretrained=False, **kwargs) -> Beit:
     model_args = dict(
         patch_size=16, embed_dim=1024, depth=24, num_heads=16,
         use_abs_pos_emb=False, use_rel_pos_bias=True, init_values=1e-5)
@@ -542,7 +542,7 @@ def beit_large_patch16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def beit_large_patch16_384(pretrained=False, **kwargs):
+def beit_large_patch16_384(pretrained=False, **kwargs) -> Beit:
     model_args = dict(
         img_size=384, patch_size=16, embed_dim=1024, depth=24, num_heads=16,
         use_abs_pos_emb=False, use_rel_pos_bias=True, init_values=1e-5)
@@ -551,7 +551,7 @@ def beit_large_patch16_384(pretrained=False, **kwargs):
 
 
 @register_model
-def beit_large_patch16_512(pretrained=False, **kwargs):
+def beit_large_patch16_512(pretrained=False, **kwargs) -> Beit:
     model_args = dict(
         img_size=512, patch_size=16, embed_dim=1024, depth=24, num_heads=16,
         use_abs_pos_emb=False, use_rel_pos_bias=True, init_values=1e-5)
@@ -560,7 +560,7 @@ def beit_large_patch16_512(pretrained=False, **kwargs):
 
 
 @register_model
-def beitv2_base_patch16_224(pretrained=False, **kwargs):
+def beitv2_base_patch16_224(pretrained=False, **kwargs) -> Beit:
     model_args = dict(
         patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4,
         use_abs_pos_emb=False, use_rel_pos_bias=True, init_values=1e-5)
@@ -569,7 +569,7 @@ def beitv2_base_patch16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def beitv2_large_patch16_224(pretrained=False, **kwargs):
+def beitv2_large_patch16_224(pretrained=False, **kwargs) -> Beit:
     model_args = dict(
         patch_size=16, embed_dim=1024, depth=24, num_heads=16,
         use_abs_pos_emb=False, use_rel_pos_bias=True, init_values=1e-5)
diff --git a/timm/models/byoanet.py b/timm/models/byoanet.py
index fbf4eefc1f..4a37ac04f5 100644
--- a/timm/models/byoanet.py
+++ b/timm/models/byoanet.py
@@ -333,7 +333,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def botnet26t_256(pretrained=False, **kwargs):
+def botnet26t_256(pretrained=False, **kwargs) -> ByobNet:
     """ Bottleneck Transformer w/ ResNet26-T backbone.
     """
     kwargs.setdefault('img_size', 256)
@@ -341,14 +341,14 @@ def botnet26t_256(pretrained=False, **kwargs):
 
 
 @register_model
-def sebotnet33ts_256(pretrained=False, **kwargs):
+def sebotnet33ts_256(pretrained=False, **kwargs) -> ByobNet:
     """ Bottleneck Transformer w/ a ResNet33-t backbone, SE attn for non Halo blocks, SiLU,
     """
     return _create_byoanet('sebotnet33ts_256', 'sebotnet33ts', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def botnet50ts_256(pretrained=False, **kwargs):
+def botnet50ts_256(pretrained=False, **kwargs) -> ByobNet:
     """ Bottleneck Transformer w/ ResNet50-T backbone, silu act.
     """
     kwargs.setdefault('img_size', 256)
@@ -356,7 +356,7 @@ def botnet50ts_256(pretrained=False, **kwargs):
 
 
 @register_model
-def eca_botnext26ts_256(pretrained=False, **kwargs):
+def eca_botnext26ts_256(pretrained=False, **kwargs) -> ByobNet:
     """ Bottleneck Transformer w/ ResNet26-T backbone, silu act.
     """
     kwargs.setdefault('img_size', 256)
@@ -364,7 +364,7 @@ def eca_botnext26ts_256(pretrained=False, **kwargs):
 
 
 @register_model
-def halonet_h1(pretrained=False, **kwargs):
+def halonet_h1(pretrained=False, **kwargs) -> ByobNet:
     """ HaloNet-H1. Halo attention in all stages as per the paper.
     NOTE: This runs very slowly!
     """
@@ -372,49 +372,49 @@ def halonet_h1(pretrained=False, **kwargs):
 
 
 @register_model
-def halonet26t(pretrained=False, **kwargs):
+def halonet26t(pretrained=False, **kwargs) -> ByobNet:
     """ HaloNet w/ a ResNet26-t backbone. Halo attention in final two stages
     """
     return _create_byoanet('halonet26t', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def sehalonet33ts(pretrained=False, **kwargs):
+def sehalonet33ts(pretrained=False, **kwargs) -> ByobNet:
     """ HaloNet w/ a ResNet33-t backbone, SE attn for non Halo blocks, SiLU, 1-2 Halo in stage 2,3,4.
     """
     return _create_byoanet('sehalonet33ts', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def halonet50ts(pretrained=False, **kwargs):
+def halonet50ts(pretrained=False, **kwargs) -> ByobNet:
     """ HaloNet w/ a ResNet50-t backbone, silu act. Halo attention in final two stages
     """
     return _create_byoanet('halonet50ts', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def eca_halonext26ts(pretrained=False, **kwargs):
+def eca_halonext26ts(pretrained=False, **kwargs) -> ByobNet:
     """ HaloNet w/ a ResNet26-t backbone, silu act. Halo attention in final two stages
     """
     return _create_byoanet('eca_halonext26ts', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def lambda_resnet26t(pretrained=False, **kwargs):
+def lambda_resnet26t(pretrained=False, **kwargs) -> ByobNet:
     """ Lambda-ResNet-26-T. Lambda layers w/ conv pos in last two stages.
     """
     return _create_byoanet('lambda_resnet26t', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def lambda_resnet50ts(pretrained=False, **kwargs):
+def lambda_resnet50ts(pretrained=False, **kwargs) -> ByobNet:
     """ Lambda-ResNet-50-TS. SiLU act. Lambda layers w/ conv pos in last two stages.
     """
     return _create_byoanet('lambda_resnet50ts', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def lambda_resnet26rpt_256(pretrained=False, **kwargs):
+def lambda_resnet26rpt_256(pretrained=False, **kwargs) -> ByobNet:
     """ Lambda-ResNet-26-R-T. Lambda layers w/ rel pos embed in last two stages.
     """
     kwargs.setdefault('img_size', 256)
@@ -422,21 +422,21 @@ def lambda_resnet26rpt_256(pretrained=False, **kwargs):
 
 
 @register_model
-def haloregnetz_b(pretrained=False, **kwargs):
+def haloregnetz_b(pretrained=False, **kwargs) -> ByobNet:
     """ Halo + RegNetZ
     """
     return _create_byoanet('haloregnetz_b', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def lamhalobotnet50ts_256(pretrained=False, **kwargs):
+def lamhalobotnet50ts_256(pretrained=False, **kwargs) -> ByobNet:
     """ Combo Attention (Lambda + Halo + Bot) Network
     """
     return _create_byoanet('lamhalobotnet50ts_256', 'lamhalobotnet50ts', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def halo2botnet50ts_256(pretrained=False, **kwargs):
+def halo2botnet50ts_256(pretrained=False, **kwargs) -> ByobNet:
     """ Combo Attention (Halo + Halo + Bot) Network
     """
     return _create_byoanet('halo2botnet50ts_256', 'halo2botnet50ts', pretrained=pretrained, **kwargs)
diff --git a/timm/models/byobnet.py b/timm/models/byobnet.py
index 50999e4a06..afb5df4282 100644
--- a/timm/models/byobnet.py
+++ b/timm/models/byobnet.py
@@ -1543,7 +1543,7 @@ def _cfgr(url='', **kwargs):
 
 
 @register_model
-def gernet_l(pretrained=False, **kwargs):
+def gernet_l(pretrained=False, **kwargs) -> ByobNet:
     """ GEResNet-Large (GENet-Large from official impl)
     `Neural Architecture Design for GPU-Efficient Networks` - https://arxiv.org/abs/2006.14090
     """
@@ -1551,7 +1551,7 @@ def gernet_l(pretrained=False, **kwargs):
 
 
 @register_model
-def gernet_m(pretrained=False, **kwargs):
+def gernet_m(pretrained=False, **kwargs) -> ByobNet:
     """ GEResNet-Medium (GENet-Normal from official impl)
     `Neural Architecture Design for GPU-Efficient Networks` - https://arxiv.org/abs/2006.14090
     """
@@ -1559,7 +1559,7 @@ def gernet_m(pretrained=False, **kwargs):
 
 
 @register_model
-def gernet_s(pretrained=False, **kwargs):
+def gernet_s(pretrained=False, **kwargs) -> ByobNet:
     """ EResNet-Small (GENet-Small from official impl)
     `Neural Architecture Design for GPU-Efficient Networks` - https://arxiv.org/abs/2006.14090
     """
@@ -1567,7 +1567,7 @@ def gernet_s(pretrained=False, **kwargs):
 
 
 @register_model
-def repvgg_a2(pretrained=False, **kwargs):
+def repvgg_a2(pretrained=False, **kwargs) -> ByobNet:
     """ RepVGG-A2
     `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
     """
@@ -1575,7 +1575,7 @@ def repvgg_a2(pretrained=False, **kwargs):
 
 
 @register_model
-def repvgg_b0(pretrained=False, **kwargs):
+def repvgg_b0(pretrained=False, **kwargs) -> ByobNet:
     """ RepVGG-B0
     `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
     """
@@ -1583,7 +1583,7 @@ def repvgg_b0(pretrained=False, **kwargs):
 
 
 @register_model
-def repvgg_b1(pretrained=False, **kwargs):
+def repvgg_b1(pretrained=False, **kwargs) -> ByobNet:
     """ RepVGG-B1
     `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
     """
@@ -1591,7 +1591,7 @@ def repvgg_b1(pretrained=False, **kwargs):
 
 
 @register_model
-def repvgg_b1g4(pretrained=False, **kwargs):
+def repvgg_b1g4(pretrained=False, **kwargs) -> ByobNet:
     """ RepVGG-B1g4
     `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
     """
@@ -1599,7 +1599,7 @@ def repvgg_b1g4(pretrained=False, **kwargs):
 
 
 @register_model
-def repvgg_b2(pretrained=False, **kwargs):
+def repvgg_b2(pretrained=False, **kwargs) -> ByobNet:
     """ RepVGG-B2
     `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
     """
@@ -1607,7 +1607,7 @@ def repvgg_b2(pretrained=False, **kwargs):
 
 
 @register_model
-def repvgg_b2g4(pretrained=False, **kwargs):
+def repvgg_b2g4(pretrained=False, **kwargs) -> ByobNet:
     """ RepVGG-B2g4
     `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
     """
@@ -1615,7 +1615,7 @@ def repvgg_b2g4(pretrained=False, **kwargs):
 
 
 @register_model
-def repvgg_b3(pretrained=False, **kwargs):
+def repvgg_b3(pretrained=False, **kwargs) -> ByobNet:
     """ RepVGG-B3
     `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
     """
@@ -1623,7 +1623,7 @@ def repvgg_b3(pretrained=False, **kwargs):
 
 
 @register_model
-def repvgg_b3g4(pretrained=False, **kwargs):
+def repvgg_b3g4(pretrained=False, **kwargs) -> ByobNet:
     """ RepVGG-B3g4
     `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
     """
@@ -1631,154 +1631,154 @@ def repvgg_b3g4(pretrained=False, **kwargs):
 
 
 @register_model
-def resnet51q(pretrained=False, **kwargs):
+def resnet51q(pretrained=False, **kwargs) -> ByobNet:
     """
     """
     return _create_byobnet('resnet51q', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def resnet61q(pretrained=False, **kwargs):
+def resnet61q(pretrained=False, **kwargs) -> ByobNet:
     """
     """
     return _create_byobnet('resnet61q', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def resnext26ts(pretrained=False, **kwargs):
+def resnext26ts(pretrained=False, **kwargs) -> ByobNet:
     """
     """
     return _create_byobnet('resnext26ts', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def gcresnext26ts(pretrained=False, **kwargs):
+def gcresnext26ts(pretrained=False, **kwargs) -> ByobNet:
     """
     """
     return _create_byobnet('gcresnext26ts', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def seresnext26ts(pretrained=False, **kwargs):
+def seresnext26ts(pretrained=False, **kwargs) -> ByobNet:
     """
     """
     return _create_byobnet('seresnext26ts', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def eca_resnext26ts(pretrained=False, **kwargs):
+def eca_resnext26ts(pretrained=False, **kwargs) -> ByobNet:
     """
     """
     return _create_byobnet('eca_resnext26ts', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def bat_resnext26ts(pretrained=False, **kwargs):
+def bat_resnext26ts(pretrained=False, **kwargs) -> ByobNet:
     """
     """
     return _create_byobnet('bat_resnext26ts', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def resnet32ts(pretrained=False, **kwargs):
+def resnet32ts(pretrained=False, **kwargs) -> ByobNet:
     """
     """
     return _create_byobnet('resnet32ts', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def resnet33ts(pretrained=False, **kwargs):
+def resnet33ts(pretrained=False, **kwargs) -> ByobNet:
     """
     """
     return _create_byobnet('resnet33ts', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def gcresnet33ts(pretrained=False, **kwargs):
+def gcresnet33ts(pretrained=False, **kwargs) -> ByobNet:
     """
     """
     return _create_byobnet('gcresnet33ts', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def seresnet33ts(pretrained=False, **kwargs):
+def seresnet33ts(pretrained=False, **kwargs) -> ByobNet:
     """
     """
     return _create_byobnet('seresnet33ts', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def eca_resnet33ts(pretrained=False, **kwargs):
+def eca_resnet33ts(pretrained=False, **kwargs) -> ByobNet:
     """
     """
     return _create_byobnet('eca_resnet33ts', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def gcresnet50t(pretrained=False, **kwargs):
+def gcresnet50t(pretrained=False, **kwargs) -> ByobNet:
     """
     """
     return _create_byobnet('gcresnet50t', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def gcresnext50ts(pretrained=False, **kwargs):
+def gcresnext50ts(pretrained=False, **kwargs) -> ByobNet:
     """
     """
     return _create_byobnet('gcresnext50ts', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def regnetz_b16(pretrained=False, **kwargs):
+def regnetz_b16(pretrained=False, **kwargs) -> ByobNet:
     """
     """
     return _create_byobnet('regnetz_b16', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def regnetz_c16(pretrained=False, **kwargs):
+def regnetz_c16(pretrained=False, **kwargs) -> ByobNet:
     """
     """
     return _create_byobnet('regnetz_c16', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def regnetz_d32(pretrained=False, **kwargs):
+def regnetz_d32(pretrained=False, **kwargs) -> ByobNet:
     """
     """
     return _create_byobnet('regnetz_d32', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def regnetz_d8(pretrained=False, **kwargs):
+def regnetz_d8(pretrained=False, **kwargs) -> ByobNet:
     """
     """
     return _create_byobnet('regnetz_d8', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def regnetz_e8(pretrained=False, **kwargs):
+def regnetz_e8(pretrained=False, **kwargs) -> ByobNet:
     """
     """
     return _create_byobnet('regnetz_e8', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def regnetz_b16_evos(pretrained=False, **kwargs):
+def regnetz_b16_evos(pretrained=False, **kwargs) -> ByobNet:
     """
     """
     return _create_byobnet('regnetz_b16_evos', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def regnetz_c16_evos(pretrained=False, **kwargs):
+def regnetz_c16_evos(pretrained=False, **kwargs) -> ByobNet:
     """
     """
     return _create_byobnet('regnetz_c16_evos', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def regnetz_d8_evos(pretrained=False, **kwargs):
+def regnetz_d8_evos(pretrained=False, **kwargs) -> ByobNet:
     """
     """
     return _create_byobnet('regnetz_d8_evos', pretrained=pretrained, **kwargs)
diff --git a/timm/models/cait.py b/timm/models/cait.py
index acd5a18480..4bc7dafc53 100644
--- a/timm/models/cait.py
+++ b/timm/models/cait.py
@@ -446,70 +446,70 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def cait_xxs24_224(pretrained=False, **kwargs):
+def cait_xxs24_224(pretrained=False, **kwargs) -> Cait:
     model_args = dict(patch_size=16, embed_dim=192, depth=24, num_heads=4, init_values=1e-5)
     model = _create_cait('cait_xxs24_224', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def cait_xxs24_384(pretrained=False, **kwargs):
+def cait_xxs24_384(pretrained=False, **kwargs) -> Cait:
     model_args = dict(patch_size=16, embed_dim=192, depth=24, num_heads=4, init_values=1e-5)
     model = _create_cait('cait_xxs24_384', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def cait_xxs36_224(pretrained=False, **kwargs):
+def cait_xxs36_224(pretrained=False, **kwargs) -> Cait:
     model_args = dict(patch_size=16, embed_dim=192, depth=36, num_heads=4, init_values=1e-5)
     model = _create_cait('cait_xxs36_224', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def cait_xxs36_384(pretrained=False, **kwargs):
+def cait_xxs36_384(pretrained=False, **kwargs) -> Cait:
     model_args = dict(patch_size=16, embed_dim=192, depth=36, num_heads=4, init_values=1e-5)
     model = _create_cait('cait_xxs36_384', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def cait_xs24_384(pretrained=False, **kwargs):
+def cait_xs24_384(pretrained=False, **kwargs) -> Cait:
     model_args = dict(patch_size=16, embed_dim=288, depth=24, num_heads=6, init_values=1e-5)
     model = _create_cait('cait_xs24_384', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def cait_s24_224(pretrained=False, **kwargs):
+def cait_s24_224(pretrained=False, **kwargs) -> Cait:
     model_args = dict(patch_size=16, embed_dim=384, depth=24, num_heads=8, init_values=1e-5)
     model = _create_cait('cait_s24_224', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def cait_s24_384(pretrained=False, **kwargs):
+def cait_s24_384(pretrained=False, **kwargs) -> Cait:
     model_args = dict(patch_size=16, embed_dim=384, depth=24, num_heads=8, init_values=1e-5)
     model = _create_cait('cait_s24_384', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def cait_s36_384(pretrained=False, **kwargs):
+def cait_s36_384(pretrained=False, **kwargs) -> Cait:
     model_args = dict(patch_size=16, embed_dim=384, depth=36, num_heads=8, init_values=1e-6)
     model = _create_cait('cait_s36_384', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def cait_m36_384(pretrained=False, **kwargs):
+def cait_m36_384(pretrained=False, **kwargs) -> Cait:
     model_args = dict(patch_size=16, embed_dim=768, depth=36, num_heads=16, init_values=1e-6)
     model = _create_cait('cait_m36_384', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def cait_m48_448(pretrained=False, **kwargs):
+def cait_m48_448(pretrained=False, **kwargs) -> Cait:
     model_args = dict(patch_size=16, embed_dim=768, depth=48, num_heads=16, init_values=1e-6)
     model = _create_cait('cait_m48_448', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
diff --git a/timm/models/coat.py b/timm/models/coat.py
index b8afbb294e..f6c89af24d 100644
--- a/timm/models/coat.py
+++ b/timm/models/coat.py
@@ -738,7 +738,7 @@ def _cfg_coat(url='', **kwargs):
 
 
 @register_model
-def coat_tiny(pretrained=False, **kwargs):
+def coat_tiny(pretrained=False, **kwargs) -> CoaT:
     model_cfg = dict(
         patch_size=4, embed_dims=[152, 152, 152, 152], serial_depths=[2, 2, 2, 2], parallel_depth=6)
     model = _create_coat('coat_tiny', pretrained=pretrained, **dict(model_cfg, **kwargs))
@@ -746,7 +746,7 @@ def coat_tiny(pretrained=False, **kwargs):
 
 
 @register_model
-def coat_mini(pretrained=False, **kwargs):
+def coat_mini(pretrained=False, **kwargs) -> CoaT:
     model_cfg = dict(
         patch_size=4, embed_dims=[152, 216, 216, 216], serial_depths=[2, 2, 2, 2], parallel_depth=6)
     model = _create_coat('coat_mini', pretrained=pretrained, **dict(model_cfg, **kwargs))
@@ -754,7 +754,7 @@ def coat_mini(pretrained=False, **kwargs):
 
 
 @register_model
-def coat_small(pretrained=False, **kwargs):
+def coat_small(pretrained=False, **kwargs) -> CoaT:
     model_cfg = dict(
         patch_size=4, embed_dims=[152, 320, 320, 320], serial_depths=[2, 2, 2, 2], parallel_depth=6, **kwargs)
     model = _create_coat('coat_small', pretrained=pretrained, **dict(model_cfg, **kwargs))
@@ -762,7 +762,7 @@ def coat_small(pretrained=False, **kwargs):
 
 
 @register_model
-def coat_lite_tiny(pretrained=False, **kwargs):
+def coat_lite_tiny(pretrained=False, **kwargs) -> CoaT:
     model_cfg = dict(
         patch_size=4, embed_dims=[64, 128, 256, 320], serial_depths=[2, 2, 2, 2], mlp_ratios=[8, 8, 4, 4])
     model = _create_coat('coat_lite_tiny', pretrained=pretrained, **dict(model_cfg, **kwargs))
@@ -770,7 +770,7 @@ def coat_lite_tiny(pretrained=False, **kwargs):
 
 
 @register_model
-def coat_lite_mini(pretrained=False, **kwargs):
+def coat_lite_mini(pretrained=False, **kwargs) -> CoaT:
     model_cfg = dict(
         patch_size=4, embed_dims=[64, 128, 320, 512], serial_depths=[2, 2, 2, 2], mlp_ratios=[8, 8, 4, 4])
     model = _create_coat('coat_lite_mini', pretrained=pretrained, **dict(model_cfg, **kwargs))
@@ -778,7 +778,7 @@ def coat_lite_mini(pretrained=False, **kwargs):
 
 
 @register_model
-def coat_lite_small(pretrained=False, **kwargs):
+def coat_lite_small(pretrained=False, **kwargs) -> CoaT:
     model_cfg = dict(
         patch_size=4, embed_dims=[64, 128, 320, 512], serial_depths=[3, 4, 6, 3], mlp_ratios=[8, 8, 4, 4])
     model = _create_coat('coat_lite_small', pretrained=pretrained, **dict(model_cfg, **kwargs))
@@ -786,7 +786,7 @@ def coat_lite_small(pretrained=False, **kwargs):
 
 
 @register_model
-def coat_lite_medium(pretrained=False, **kwargs):
+def coat_lite_medium(pretrained=False, **kwargs) -> CoaT:
     model_cfg = dict(
         patch_size=4, embed_dims=[128, 256, 320, 512], serial_depths=[3, 6, 10, 8])
     model = _create_coat('coat_lite_medium', pretrained=pretrained, **dict(model_cfg, **kwargs))
@@ -794,7 +794,7 @@ def coat_lite_medium(pretrained=False, **kwargs):
 
 
 @register_model
-def coat_lite_medium_384(pretrained=False, **kwargs):
+def coat_lite_medium_384(pretrained=False, **kwargs) -> CoaT:
     model_cfg = dict(
         img_size=384, patch_size=4, embed_dims=[128, 256, 320, 512], serial_depths=[3, 6, 10, 8])
     model = _create_coat('coat_lite_medium_384', pretrained=pretrained, **dict(model_cfg, **kwargs))
diff --git a/timm/models/convit.py b/timm/models/convit.py
index 16e8b5c4e9..6cfcae2705 100644
--- a/timm/models/convit.py
+++ b/timm/models/convit.py
@@ -35,7 +35,7 @@
 from .vision_transformer_hybrid import HybridEmbed
 
 
-__all__ = ['ConViT']
+__all__ = ['ConVit']
 
 
 @register_notrace_module  # reason: FX can't symbolically trace control flow in forward method
@@ -237,7 +237,7 @@ def forward(self, x):
         return x
 
 
-class ConViT(nn.Module):
+class ConVit(nn.Module):
     """ Vision Transformer with support for patch or hybrid CNN input stage
     """
 
@@ -385,7 +385,7 @@ def _create_convit(variant, pretrained=False, **kwargs):
     if kwargs.get('features_only', None):
         raise RuntimeError('features_only not implemented for Vision Transformer models.')
 
-    return build_model_with_cfg(ConViT, variant, pretrained, **kwargs)
+    return build_model_with_cfg(ConVit, variant, pretrained, **kwargs)
 
 
 def _cfg(url='', **kwargs):
@@ -407,7 +407,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def convit_tiny(pretrained=False, **kwargs):
+def convit_tiny(pretrained=False, **kwargs) -> ConVit:
     model_args = dict(
         local_up_to_layer=10, locality_strength=1.0, embed_dim=48, num_heads=4)
     model = _create_convit(variant='convit_tiny', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -415,7 +415,7 @@ def convit_tiny(pretrained=False, **kwargs):
 
 
 @register_model
-def convit_small(pretrained=False, **kwargs):
+def convit_small(pretrained=False, **kwargs) -> ConVit:
     model_args = dict(
         local_up_to_layer=10, locality_strength=1.0, embed_dim=48, num_heads=9)
     model = _create_convit(variant='convit_small', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -423,7 +423,7 @@ def convit_small(pretrained=False, **kwargs):
 
 
 @register_model
-def convit_base(pretrained=False, **kwargs):
+def convit_base(pretrained=False, **kwargs) -> ConVit:
     model_args = dict(
         local_up_to_layer=10, locality_strength=1.0, embed_dim=48, num_heads=16)
     model = _create_convit(variant='convit_base', pretrained=pretrained, **dict(model_args, **kwargs))
diff --git a/timm/models/convmixer.py b/timm/models/convmixer.py
index edb3c3aebd..55220bc084 100644
--- a/timm/models/convmixer.py
+++ b/timm/models/convmixer.py
@@ -124,18 +124,18 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def convmixer_1536_20(pretrained=False, **kwargs):
+def convmixer_1536_20(pretrained=False, **kwargs) -> ConvMixer:
     model_args = dict(dim=1536, depth=20, kernel_size=9, patch_size=7, **kwargs)
     return _create_convmixer('convmixer_1536_20', pretrained, **model_args)
 
 
 @register_model
-def convmixer_768_32(pretrained=False, **kwargs):
+def convmixer_768_32(pretrained=False, **kwargs) -> ConvMixer:
     model_args = dict(dim=768, depth=32, kernel_size=7, patch_size=7, act_layer=nn.ReLU, **kwargs)
     return _create_convmixer('convmixer_768_32', pretrained, **model_args)
 
 
 @register_model
-def convmixer_1024_20_ks9_p14(pretrained=False, **kwargs):
+def convmixer_1024_20_ks9_p14(pretrained=False, **kwargs) -> ConvMixer:
     model_args = dict(dim=1024, depth=20, kernel_size=9, patch_size=14, **kwargs)
     return _create_convmixer('convmixer_1024_20_ks9_p14', pretrained, **model_args)
\ No newline at end of file
diff --git a/timm/models/convnext.py b/timm/models/convnext.py
index 8e2c3e24e0..80b207e75d 100644
--- a/timm/models/convnext.py
+++ b/timm/models/convnext.py
@@ -884,7 +884,7 @@ def _cfgv2(url='', **kwargs):
 
 
 @register_model
-def convnext_atto(pretrained=False, **kwargs):
+def convnext_atto(pretrained=False, **kwargs) -> ConvNeXt:
     # timm femto variant (NOTE: still tweaking depths, will vary between 3-4M param, current is 3.7M
     model_args = dict(depths=(2, 2, 6, 2), dims=(40, 80, 160, 320), conv_mlp=True)
     model = _create_convnext('convnext_atto', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -892,7 +892,7 @@ def convnext_atto(pretrained=False, **kwargs):
 
 
 @register_model
-def convnext_atto_ols(pretrained=False, **kwargs):
+def convnext_atto_ols(pretrained=False, **kwargs) -> ConvNeXt:
     # timm femto variant with overlapping 3x3 conv stem, wider than non-ols femto above, current param count 3.7M
     model_args = dict(depths=(2, 2, 6, 2), dims=(40, 80, 160, 320), conv_mlp=True, stem_type='overlap_tiered')
     model = _create_convnext('convnext_atto_ols', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -900,7 +900,7 @@ def convnext_atto_ols(pretrained=False, **kwargs):
 
 
 @register_model
-def convnext_femto(pretrained=False, **kwargs):
+def convnext_femto(pretrained=False, **kwargs) -> ConvNeXt:
     # timm femto variant
     model_args = dict(depths=(2, 2, 6, 2), dims=(48, 96, 192, 384), conv_mlp=True)
     model = _create_convnext('convnext_femto', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -908,7 +908,7 @@ def convnext_femto(pretrained=False, **kwargs):
 
 
 @register_model
-def convnext_femto_ols(pretrained=False, **kwargs):
+def convnext_femto_ols(pretrained=False, **kwargs) -> ConvNeXt:
     # timm femto variant
     model_args = dict(depths=(2, 2, 6, 2), dims=(48, 96, 192, 384), conv_mlp=True, stem_type='overlap_tiered')
     model = _create_convnext('convnext_femto_ols', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -916,7 +916,7 @@ def convnext_femto_ols(pretrained=False, **kwargs):
 
 
 @register_model
-def convnext_pico(pretrained=False, **kwargs):
+def convnext_pico(pretrained=False, **kwargs) -> ConvNeXt:
     # timm pico variant
     model_args = dict(depths=(2, 2, 6, 2), dims=(64, 128, 256, 512), conv_mlp=True)
     model = _create_convnext('convnext_pico', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -924,7 +924,7 @@ def convnext_pico(pretrained=False, **kwargs):
 
 
 @register_model
-def convnext_pico_ols(pretrained=False, **kwargs):
+def convnext_pico_ols(pretrained=False, **kwargs) -> ConvNeXt:
     # timm nano variant with overlapping 3x3 conv stem
     model_args = dict(depths=(2, 2, 6, 2), dims=(64, 128, 256, 512), conv_mlp=True,  stem_type='overlap_tiered')
     model = _create_convnext('convnext_pico_ols', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -932,7 +932,7 @@ def convnext_pico_ols(pretrained=False, **kwargs):
 
 
 @register_model
-def convnext_nano(pretrained=False, **kwargs):
+def convnext_nano(pretrained=False, **kwargs) -> ConvNeXt:
     # timm nano variant with standard stem and head
     model_args = dict(depths=(2, 2, 8, 2), dims=(80, 160, 320, 640), conv_mlp=True)
     model = _create_convnext('convnext_nano', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -940,7 +940,7 @@ def convnext_nano(pretrained=False, **kwargs):
 
 
 @register_model
-def convnext_nano_ols(pretrained=False, **kwargs):
+def convnext_nano_ols(pretrained=False, **kwargs) -> ConvNeXt:
     # experimental nano variant with overlapping conv stem
     model_args = dict(depths=(2, 2, 8, 2), dims=(80, 160, 320, 640), conv_mlp=True, stem_type='overlap')
     model = _create_convnext('convnext_nano_ols', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -948,7 +948,7 @@ def convnext_nano_ols(pretrained=False, **kwargs):
 
 
 @register_model
-def convnext_tiny_hnf(pretrained=False, **kwargs):
+def convnext_tiny_hnf(pretrained=False, **kwargs) -> ConvNeXt:
     # experimental tiny variant with norm before pooling in head (head norm first)
     model_args = dict(depths=(3, 3, 9, 3), dims=(96, 192, 384, 768), head_norm_first=True, conv_mlp=True)
     model = _create_convnext('convnext_tiny_hnf', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -956,56 +956,56 @@ def convnext_tiny_hnf(pretrained=False, **kwargs):
 
 
 @register_model
-def convnext_tiny(pretrained=False, **kwargs):
+def convnext_tiny(pretrained=False, **kwargs) -> ConvNeXt:
     model_args = dict(depths=(3, 3, 9, 3), dims=(96, 192, 384, 768))
     model = _create_convnext('convnext_tiny', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def convnext_small(pretrained=False, **kwargs):
+def convnext_small(pretrained=False, **kwargs) -> ConvNeXt:
     model_args = dict(depths=[3, 3, 27, 3], dims=[96, 192, 384, 768])
     model = _create_convnext('convnext_small', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def convnext_base(pretrained=False, **kwargs):
+def convnext_base(pretrained=False, **kwargs) -> ConvNeXt:
     model_args = dict(depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024])
     model = _create_convnext('convnext_base', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def convnext_large(pretrained=False, **kwargs):
+def convnext_large(pretrained=False, **kwargs) -> ConvNeXt:
     model_args = dict(depths=[3, 3, 27, 3], dims=[192, 384, 768, 1536])
     model = _create_convnext('convnext_large', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def convnext_large_mlp(pretrained=False, **kwargs):
+def convnext_large_mlp(pretrained=False, **kwargs) -> ConvNeXt:
     model_args = dict(depths=[3, 3, 27, 3], dims=[192, 384, 768, 1536], head_hidden_size=1536)
     model = _create_convnext('convnext_large_mlp', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def convnext_xlarge(pretrained=False, **kwargs):
+def convnext_xlarge(pretrained=False, **kwargs) -> ConvNeXt:
     model_args = dict(depths=[3, 3, 27, 3], dims=[256, 512, 1024, 2048])
     model = _create_convnext('convnext_xlarge', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def convnext_xxlarge(pretrained=False, **kwargs):
+def convnext_xxlarge(pretrained=False, **kwargs) -> ConvNeXt:
     model_args = dict(depths=[3, 4, 30, 3], dims=[384, 768, 1536, 3072], norm_eps=kwargs.pop('norm_eps', 1e-5))
     model = _create_convnext('convnext_xxlarge', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def convnextv2_atto(pretrained=False, **kwargs):
+def convnextv2_atto(pretrained=False, **kwargs) -> ConvNeXt:
     # timm femto variant (NOTE: still tweaking depths, will vary between 3-4M param, current is 3.7M
     model_args = dict(
         depths=(2, 2, 6, 2), dims=(40, 80, 160, 320), use_grn=True, ls_init_value=None, conv_mlp=True)
@@ -1014,7 +1014,7 @@ def convnextv2_atto(pretrained=False, **kwargs):
 
 
 @register_model
-def convnextv2_femto(pretrained=False, **kwargs):
+def convnextv2_femto(pretrained=False, **kwargs) -> ConvNeXt:
     # timm femto variant
     model_args = dict(
         depths=(2, 2, 6, 2), dims=(48, 96, 192, 384), use_grn=True, ls_init_value=None, conv_mlp=True)
@@ -1023,7 +1023,7 @@ def convnextv2_femto(pretrained=False, **kwargs):
 
 
 @register_model
-def convnextv2_pico(pretrained=False, **kwargs):
+def convnextv2_pico(pretrained=False, **kwargs) -> ConvNeXt:
     # timm pico variant
     model_args = dict(
         depths=(2, 2, 6, 2), dims=(64, 128, 256, 512), use_grn=True, ls_init_value=None, conv_mlp=True)
@@ -1032,7 +1032,7 @@ def convnextv2_pico(pretrained=False, **kwargs):
 
 
 @register_model
-def convnextv2_nano(pretrained=False, **kwargs):
+def convnextv2_nano(pretrained=False, **kwargs) -> ConvNeXt:
     # timm nano variant with standard stem and head
     model_args = dict(
         depths=(2, 2, 8, 2), dims=(80, 160, 320, 640), use_grn=True, ls_init_value=None, conv_mlp=True)
@@ -1041,35 +1041,35 @@ def convnextv2_nano(pretrained=False, **kwargs):
 
 
 @register_model
-def convnextv2_tiny(pretrained=False, **kwargs):
+def convnextv2_tiny(pretrained=False, **kwargs) -> ConvNeXt:
     model_args = dict(depths=(3, 3, 9, 3), dims=(96, 192, 384, 768), use_grn=True, ls_init_value=None)
     model = _create_convnext('convnextv2_tiny', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def convnextv2_small(pretrained=False, **kwargs):
+def convnextv2_small(pretrained=False, **kwargs) -> ConvNeXt:
     model_args = dict(depths=[3, 3, 27, 3], dims=[96, 192, 384, 768], use_grn=True, ls_init_value=None)
     model = _create_convnext('convnextv2_small', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def convnextv2_base(pretrained=False, **kwargs):
+def convnextv2_base(pretrained=False, **kwargs) -> ConvNeXt:
     model_args = dict(depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024], use_grn=True, ls_init_value=None)
     model = _create_convnext('convnextv2_base', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def convnextv2_large(pretrained=False, **kwargs):
+def convnextv2_large(pretrained=False, **kwargs) -> ConvNeXt:
     model_args = dict(depths=[3, 3, 27, 3], dims=[192, 384, 768, 1536], use_grn=True, ls_init_value=None)
     model = _create_convnext('convnextv2_large', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def convnextv2_huge(pretrained=False, **kwargs):
+def convnextv2_huge(pretrained=False, **kwargs) -> ConvNeXt:
     model_args = dict(depths=[3, 3, 27, 3], dims=[352, 704, 1408, 2816], use_grn=True, ls_init_value=None)
     model = _create_convnext('convnextv2_huge', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
diff --git a/timm/models/crossvit.py b/timm/models/crossvit.py
index 39709e4426..5c90aec97c 100644
--- a/timm/models/crossvit.py
+++ b/timm/models/crossvit.py
@@ -39,7 +39,7 @@
 from ._registry import register_model, generate_default_cfgs
 from .vision_transformer import Block
 
-__all__ = ['CrossViT']  # model_registry will add each entrypoint fn to this
+__all__ = ['CrossVit']  # model_registry will add each entrypoint fn to this
 
 
 class PatchEmbed(nn.Module):
@@ -293,7 +293,7 @@ def scale_image(x, ss: Tuple[int, int], crop_scale: bool = False):  # annotation
     return x
 
 
-class CrossViT(nn.Module):
+class CrossVit(nn.Module):
     """ Vision Transformer with support for patch or hybrid CNN input stage
     """
 
@@ -479,7 +479,7 @@ def pretrained_filter_fn(state_dict):
         return new_state_dict
 
     return build_model_with_cfg(
-        CrossViT,
+        CrossVit,
         variant,
         pretrained,
         pretrained_filter_fn=pretrained_filter_fn,
@@ -529,7 +529,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def crossvit_tiny_240(pretrained=False, **kwargs):
+def crossvit_tiny_240(pretrained=False, **kwargs) -> CrossVit:
     model_args = dict(
         img_scale=(1.0, 224/240), patch_size=[12, 16], embed_dim=[96, 192], depth=[[1, 4, 0], [1, 4, 0], [1, 4, 0]],
         num_heads=[3, 3], mlp_ratio=[4, 4, 1])
@@ -538,7 +538,7 @@ def crossvit_tiny_240(pretrained=False, **kwargs):
 
 
 @register_model
-def crossvit_small_240(pretrained=False, **kwargs):
+def crossvit_small_240(pretrained=False, **kwargs) -> CrossVit:
     model_args = dict(
         img_scale=(1.0, 224/240), patch_size=[12, 16], embed_dim=[192, 384], depth=[[1, 4, 0], [1, 4, 0], [1, 4, 0]],
         num_heads=[6, 6], mlp_ratio=[4, 4, 1])
@@ -547,7 +547,7 @@ def crossvit_small_240(pretrained=False, **kwargs):
 
 
 @register_model
-def crossvit_base_240(pretrained=False, **kwargs):
+def crossvit_base_240(pretrained=False, **kwargs) -> CrossVit:
     model_args = dict(
         img_scale=(1.0, 224/240), patch_size=[12, 16], embed_dim=[384, 768], depth=[[1, 4, 0], [1, 4, 0], [1, 4, 0]],
         num_heads=[12, 12], mlp_ratio=[4, 4, 1])
@@ -556,7 +556,7 @@ def crossvit_base_240(pretrained=False, **kwargs):
 
 
 @register_model
-def crossvit_9_240(pretrained=False, **kwargs):
+def crossvit_9_240(pretrained=False, **kwargs) -> CrossVit:
     model_args = dict(
         img_scale=(1.0, 224/240), patch_size=[12, 16], embed_dim=[128, 256], depth=[[1, 3, 0], [1, 3, 0], [1, 3, 0]],
         num_heads=[4, 4], mlp_ratio=[3, 3, 1])
@@ -565,7 +565,7 @@ def crossvit_9_240(pretrained=False, **kwargs):
 
 
 @register_model
-def crossvit_15_240(pretrained=False, **kwargs):
+def crossvit_15_240(pretrained=False, **kwargs) -> CrossVit:
     model_args = dict(
         img_scale=(1.0, 224/240), patch_size=[12, 16], embed_dim=[192, 384], depth=[[1, 5, 0], [1, 5, 0], [1, 5, 0]],
         num_heads=[6, 6], mlp_ratio=[3, 3, 1])
@@ -574,7 +574,7 @@ def crossvit_15_240(pretrained=False, **kwargs):
 
 
 @register_model
-def crossvit_18_240(pretrained=False, **kwargs):
+def crossvit_18_240(pretrained=False, **kwargs) -> CrossVit:
     model_args = dict(
         img_scale=(1.0, 224 / 240), patch_size=[12, 16], embed_dim=[224, 448], depth=[[1, 6, 0], [1, 6, 0], [1, 6, 0]],
         num_heads=[7, 7], mlp_ratio=[3, 3, 1], **kwargs)
@@ -583,7 +583,7 @@ def crossvit_18_240(pretrained=False, **kwargs):
 
 
 @register_model
-def crossvit_9_dagger_240(pretrained=False, **kwargs):
+def crossvit_9_dagger_240(pretrained=False, **kwargs) -> CrossVit:
     model_args = dict(
         img_scale=(1.0, 224 / 240), patch_size=[12, 16], embed_dim=[128, 256], depth=[[1, 3, 0], [1, 3, 0], [1, 3, 0]],
         num_heads=[4, 4], mlp_ratio=[3, 3, 1], multi_conv=True)
@@ -592,7 +592,7 @@ def crossvit_9_dagger_240(pretrained=False, **kwargs):
 
 
 @register_model
-def crossvit_15_dagger_240(pretrained=False, **kwargs):
+def crossvit_15_dagger_240(pretrained=False, **kwargs) -> CrossVit:
     model_args = dict(
         img_scale=(1.0, 224/240), patch_size=[12, 16], embed_dim=[192, 384], depth=[[1, 5, 0], [1, 5, 0], [1, 5, 0]],
         num_heads=[6, 6], mlp_ratio=[3, 3, 1], multi_conv=True)
@@ -601,7 +601,7 @@ def crossvit_15_dagger_240(pretrained=False, **kwargs):
 
 
 @register_model
-def crossvit_15_dagger_408(pretrained=False, **kwargs):
+def crossvit_15_dagger_408(pretrained=False, **kwargs) -> CrossVit:
     model_args = dict(
         img_scale=(1.0, 384/408), patch_size=[12, 16], embed_dim=[192, 384], depth=[[1, 5, 0], [1, 5, 0], [1, 5, 0]],
         num_heads=[6, 6], mlp_ratio=[3, 3, 1], multi_conv=True)
@@ -610,7 +610,7 @@ def crossvit_15_dagger_408(pretrained=False, **kwargs):
 
 
 @register_model
-def crossvit_18_dagger_240(pretrained=False, **kwargs):
+def crossvit_18_dagger_240(pretrained=False, **kwargs) -> CrossVit:
     model_args = dict(
         img_scale=(1.0, 224/240), patch_size=[12, 16], embed_dim=[224, 448], depth=[[1, 6, 0], [1, 6, 0], [1, 6, 0]],
         num_heads=[7, 7], mlp_ratio=[3, 3, 1], multi_conv=True)
@@ -619,7 +619,7 @@ def crossvit_18_dagger_240(pretrained=False, **kwargs):
 
 
 @register_model
-def crossvit_18_dagger_408(pretrained=False, **kwargs):
+def crossvit_18_dagger_408(pretrained=False, **kwargs) -> CrossVit:
     model_args = dict(
         img_scale=(1.0, 384/408), patch_size=[12, 16], embed_dim=[224, 448], depth=[[1, 6, 0], [1, 6, 0], [1, 6, 0]],
         num_heads=[7, 7], mlp_ratio=[3, 3, 1], multi_conv=True)
diff --git a/timm/models/cspnet.py b/timm/models/cspnet.py
index 96653cbb14..21b2cd344a 100644
--- a/timm/models/cspnet.py
+++ b/timm/models/cspnet.py
@@ -132,167 +132,6 @@ def _cs3_cfg(
     )
 
 
-model_cfgs = dict(
-    cspresnet50=CspModelCfg(
-        stem=CspStemCfg(out_chs=64, kernel_size=7, stride=4, pool='max'),
-        stages=CspStagesCfg(
-            depth=(3, 3, 5, 2),
-            out_chs=(128, 256, 512, 1024),
-            stride=(1, 2),
-            expand_ratio=2.,
-            bottle_ratio=0.5,
-            cross_linear=True,
-        ),
-    ),
-    cspresnet50d=CspModelCfg(
-        stem=CspStemCfg(out_chs=(32, 32, 64), kernel_size=3, stride=4, pool='max'),
-        stages=CspStagesCfg(
-            depth=(3, 3, 5, 2),
-            out_chs=(128, 256, 512, 1024),
-            stride=(1,) + (2,),
-            expand_ratio=2.,
-            bottle_ratio=0.5,
-            block_ratio=1.,
-            cross_linear=True,
-        ),
-    ),
-    cspresnet50w=CspModelCfg(
-        stem=CspStemCfg(out_chs=(32, 32, 64), kernel_size=3, stride=4, pool='max'),
-        stages=CspStagesCfg(
-            depth=(3, 3, 5, 2),
-            out_chs=(256, 512, 1024, 2048),
-            stride=(1,) + (2,),
-            expand_ratio=1.,
-            bottle_ratio=0.25,
-            block_ratio=0.5,
-            cross_linear=True,
-        ),
-    ),
-    cspresnext50=CspModelCfg(
-        stem=CspStemCfg(out_chs=64, kernel_size=7, stride=4, pool='max'),
-        stages=CspStagesCfg(
-            depth=(3, 3, 5, 2),
-            out_chs=(256, 512, 1024, 2048),
-            stride=(1,) + (2,),
-            groups=32,
-            expand_ratio=1.,
-            bottle_ratio=1.,
-            block_ratio=0.5,
-            cross_linear=True,
-        ),
-    ),
-    cspdarknet53=CspModelCfg(
-        stem=CspStemCfg(out_chs=32, kernel_size=3, stride=1, pool=''),
-        stages=CspStagesCfg(
-            depth=(1, 2, 8, 8, 4),
-            out_chs=(64, 128, 256, 512, 1024),
-            stride=2,
-            expand_ratio=(2.,) + (1.,),
-            bottle_ratio=(0.5,) + (1.,),
-            block_ratio=(1.,) + (0.5,),
-            down_growth=True,
-            block_type='dark',
-        ),
-    ),
-    darknet17=CspModelCfg(
-        stem=CspStemCfg(out_chs=32, kernel_size=3, stride=1, pool=''),
-        stages=CspStagesCfg(
-            depth=(1,) * 5,
-            out_chs=(64, 128, 256, 512, 1024),
-            stride=(2,),
-            bottle_ratio=(0.5,),
-            block_ratio=(1.,),
-            stage_type='dark',
-            block_type='dark',
-        ),
-    ),
-    darknet21=CspModelCfg(
-        stem=CspStemCfg(out_chs=32, kernel_size=3, stride=1, pool=''),
-        stages=CspStagesCfg(
-            depth=(1, 1, 1, 2, 2),
-            out_chs=(64, 128, 256, 512, 1024),
-            stride=(2,),
-            bottle_ratio=(0.5,),
-            block_ratio=(1.,),
-            stage_type='dark',
-            block_type='dark',
-
-        ),
-    ),
-    sedarknet21=CspModelCfg(
-        stem=CspStemCfg(out_chs=32, kernel_size=3, stride=1, pool=''),
-        stages=CspStagesCfg(
-            depth=(1, 1, 1, 2, 2),
-            out_chs=(64, 128, 256, 512, 1024),
-            stride=2,
-            bottle_ratio=0.5,
-            block_ratio=1.,
-            attn_layer='se',
-            stage_type='dark',
-            block_type='dark',
-
-        ),
-    ),
-    darknet53=CspModelCfg(
-        stem=CspStemCfg(out_chs=32, kernel_size=3, stride=1, pool=''),
-        stages=CspStagesCfg(
-            depth=(1, 2, 8, 8, 4),
-            out_chs=(64, 128, 256, 512, 1024),
-            stride=2,
-            bottle_ratio=0.5,
-            block_ratio=1.,
-            stage_type='dark',
-            block_type='dark',
-        ),
-    ),
-    darknetaa53=CspModelCfg(
-        stem=CspStemCfg(out_chs=32, kernel_size=3, stride=1, pool=''),
-        stages=CspStagesCfg(
-            depth=(1, 2, 8, 8, 4),
-            out_chs=(64, 128, 256, 512, 1024),
-            stride=2,
-            bottle_ratio=0.5,
-            block_ratio=1.,
-            avg_down=True,
-            stage_type='dark',
-            block_type='dark',
-        ),
-    ),
-
-    cs3darknet_s=_cs3_cfg(width_multiplier=0.5, depth_multiplier=0.5),
-    cs3darknet_m=_cs3_cfg(width_multiplier=0.75, depth_multiplier=0.67),
-    cs3darknet_l=_cs3_cfg(),
-    cs3darknet_x=_cs3_cfg(width_multiplier=1.25, depth_multiplier=1.33),
-
-    cs3darknet_focus_s=_cs3_cfg(width_multiplier=0.5, depth_multiplier=0.5, focus=True),
-    cs3darknet_focus_m=_cs3_cfg(width_multiplier=0.75, depth_multiplier=0.67, focus=True),
-    cs3darknet_focus_l=_cs3_cfg(focus=True),
-    cs3darknet_focus_x=_cs3_cfg(width_multiplier=1.25, depth_multiplier=1.33, focus=True),
-
-    cs3sedarknet_l=_cs3_cfg(attn_layer='se', attn_kwargs=dict(rd_ratio=.25)),
-    cs3sedarknet_x=_cs3_cfg(attn_layer='se', width_multiplier=1.25, depth_multiplier=1.33),
-
-    cs3sedarknet_xdw=CspModelCfg(
-        stem=CspStemCfg(out_chs=(32, 64), kernel_size=3, stride=2, pool=''),
-        stages=CspStagesCfg(
-            depth=(3, 6, 12, 4),
-            out_chs=(256, 512, 1024, 2048),
-            stride=2,
-            groups=(1, 1, 256, 512),
-            bottle_ratio=0.5,
-            block_ratio=0.5,
-            attn_layer='se',
-        ),
-        act_layer='silu',
-    ),
-
-    cs3edgenet_x=_cs3_cfg(width_multiplier=1.25, depth_multiplier=1.33, bottle_ratio=1.5, block_type='edge'),
-    cs3se_edgenet_x=_cs3_cfg(
-        width_multiplier=1.25, depth_multiplier=1.33, bottle_ratio=1.5, block_type='edge',
-        attn_layer='se', attn_kwargs=dict(rd_ratio=.25)),
-)
-
-
 class BottleneckBlock(nn.Module):
     """ ResNe(X)t Bottleneck Block
     """
@@ -892,6 +731,167 @@ def _init_weights(module, name, zero_init_last=False):
         module.zero_init_last()
 
 
+model_cfgs = dict(
+    cspresnet50=CspModelCfg(
+        stem=CspStemCfg(out_chs=64, kernel_size=7, stride=4, pool='max'),
+        stages=CspStagesCfg(
+            depth=(3, 3, 5, 2),
+            out_chs=(128, 256, 512, 1024),
+            stride=(1, 2),
+            expand_ratio=2.,
+            bottle_ratio=0.5,
+            cross_linear=True,
+        ),
+    ),
+    cspresnet50d=CspModelCfg(
+        stem=CspStemCfg(out_chs=(32, 32, 64), kernel_size=3, stride=4, pool='max'),
+        stages=CspStagesCfg(
+            depth=(3, 3, 5, 2),
+            out_chs=(128, 256, 512, 1024),
+            stride=(1,) + (2,),
+            expand_ratio=2.,
+            bottle_ratio=0.5,
+            block_ratio=1.,
+            cross_linear=True,
+        ),
+    ),
+    cspresnet50w=CspModelCfg(
+        stem=CspStemCfg(out_chs=(32, 32, 64), kernel_size=3, stride=4, pool='max'),
+        stages=CspStagesCfg(
+            depth=(3, 3, 5, 2),
+            out_chs=(256, 512, 1024, 2048),
+            stride=(1,) + (2,),
+            expand_ratio=1.,
+            bottle_ratio=0.25,
+            block_ratio=0.5,
+            cross_linear=True,
+        ),
+    ),
+    cspresnext50=CspModelCfg(
+        stem=CspStemCfg(out_chs=64, kernel_size=7, stride=4, pool='max'),
+        stages=CspStagesCfg(
+            depth=(3, 3, 5, 2),
+            out_chs=(256, 512, 1024, 2048),
+            stride=(1,) + (2,),
+            groups=32,
+            expand_ratio=1.,
+            bottle_ratio=1.,
+            block_ratio=0.5,
+            cross_linear=True,
+        ),
+    ),
+    cspdarknet53=CspModelCfg(
+        stem=CspStemCfg(out_chs=32, kernel_size=3, stride=1, pool=''),
+        stages=CspStagesCfg(
+            depth=(1, 2, 8, 8, 4),
+            out_chs=(64, 128, 256, 512, 1024),
+            stride=2,
+            expand_ratio=(2.,) + (1.,),
+            bottle_ratio=(0.5,) + (1.,),
+            block_ratio=(1.,) + (0.5,),
+            down_growth=True,
+            block_type='dark',
+        ),
+    ),
+    darknet17=CspModelCfg(
+        stem=CspStemCfg(out_chs=32, kernel_size=3, stride=1, pool=''),
+        stages=CspStagesCfg(
+            depth=(1,) * 5,
+            out_chs=(64, 128, 256, 512, 1024),
+            stride=(2,),
+            bottle_ratio=(0.5,),
+            block_ratio=(1.,),
+            stage_type='dark',
+            block_type='dark',
+        ),
+    ),
+    darknet21=CspModelCfg(
+        stem=CspStemCfg(out_chs=32, kernel_size=3, stride=1, pool=''),
+        stages=CspStagesCfg(
+            depth=(1, 1, 1, 2, 2),
+            out_chs=(64, 128, 256, 512, 1024),
+            stride=(2,),
+            bottle_ratio=(0.5,),
+            block_ratio=(1.,),
+            stage_type='dark',
+            block_type='dark',
+
+        ),
+    ),
+    sedarknet21=CspModelCfg(
+        stem=CspStemCfg(out_chs=32, kernel_size=3, stride=1, pool=''),
+        stages=CspStagesCfg(
+            depth=(1, 1, 1, 2, 2),
+            out_chs=(64, 128, 256, 512, 1024),
+            stride=2,
+            bottle_ratio=0.5,
+            block_ratio=1.,
+            attn_layer='se',
+            stage_type='dark',
+            block_type='dark',
+
+        ),
+    ),
+    darknet53=CspModelCfg(
+        stem=CspStemCfg(out_chs=32, kernel_size=3, stride=1, pool=''),
+        stages=CspStagesCfg(
+            depth=(1, 2, 8, 8, 4),
+            out_chs=(64, 128, 256, 512, 1024),
+            stride=2,
+            bottle_ratio=0.5,
+            block_ratio=1.,
+            stage_type='dark',
+            block_type='dark',
+        ),
+    ),
+    darknetaa53=CspModelCfg(
+        stem=CspStemCfg(out_chs=32, kernel_size=3, stride=1, pool=''),
+        stages=CspStagesCfg(
+            depth=(1, 2, 8, 8, 4),
+            out_chs=(64, 128, 256, 512, 1024),
+            stride=2,
+            bottle_ratio=0.5,
+            block_ratio=1.,
+            avg_down=True,
+            stage_type='dark',
+            block_type='dark',
+        ),
+    ),
+
+    cs3darknet_s=_cs3_cfg(width_multiplier=0.5, depth_multiplier=0.5),
+    cs3darknet_m=_cs3_cfg(width_multiplier=0.75, depth_multiplier=0.67),
+    cs3darknet_l=_cs3_cfg(),
+    cs3darknet_x=_cs3_cfg(width_multiplier=1.25, depth_multiplier=1.33),
+
+    cs3darknet_focus_s=_cs3_cfg(width_multiplier=0.5, depth_multiplier=0.5, focus=True),
+    cs3darknet_focus_m=_cs3_cfg(width_multiplier=0.75, depth_multiplier=0.67, focus=True),
+    cs3darknet_focus_l=_cs3_cfg(focus=True),
+    cs3darknet_focus_x=_cs3_cfg(width_multiplier=1.25, depth_multiplier=1.33, focus=True),
+
+    cs3sedarknet_l=_cs3_cfg(attn_layer='se', attn_kwargs=dict(rd_ratio=.25)),
+    cs3sedarknet_x=_cs3_cfg(attn_layer='se', width_multiplier=1.25, depth_multiplier=1.33),
+
+    cs3sedarknet_xdw=CspModelCfg(
+        stem=CspStemCfg(out_chs=(32, 64), kernel_size=3, stride=2, pool=''),
+        stages=CspStagesCfg(
+            depth=(3, 6, 12, 4),
+            out_chs=(256, 512, 1024, 2048),
+            stride=2,
+            groups=(1, 1, 256, 512),
+            bottle_ratio=0.5,
+            block_ratio=0.5,
+            attn_layer='se',
+        ),
+        act_layer='silu',
+    ),
+
+    cs3edgenet_x=_cs3_cfg(width_multiplier=1.25, depth_multiplier=1.33, bottle_ratio=1.5, block_type='edge'),
+    cs3se_edgenet_x=_cs3_cfg(
+        width_multiplier=1.25, depth_multiplier=1.33, bottle_ratio=1.5, block_type='edge',
+        attn_layer='se', attn_kwargs=dict(rd_ratio=.25)),
+)
+
+
 def _create_cspnet(variant, pretrained=False, **kwargs):
     if variant.startswith('darknet') or variant.startswith('cspdarknet'):
         # NOTE: DarkNet is one of few models with stride==1 features w/ 6 out_indices [0..5]
@@ -992,115 +992,115 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def cspresnet50(pretrained=False, **kwargs):
+def cspresnet50(pretrained=False, **kwargs) -> CspNet:
     return _create_cspnet('cspresnet50', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def cspresnet50d(pretrained=False, **kwargs):
+def cspresnet50d(pretrained=False, **kwargs) -> CspNet:
     return _create_cspnet('cspresnet50d', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def cspresnet50w(pretrained=False, **kwargs):
+def cspresnet50w(pretrained=False, **kwargs) -> CspNet:
     return _create_cspnet('cspresnet50w', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def cspresnext50(pretrained=False, **kwargs):
+def cspresnext50(pretrained=False, **kwargs) -> CspNet:
     return _create_cspnet('cspresnext50', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def cspdarknet53(pretrained=False, **kwargs):
+def cspdarknet53(pretrained=False, **kwargs) -> CspNet:
     return _create_cspnet('cspdarknet53', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def darknet17(pretrained=False, **kwargs):
+def darknet17(pretrained=False, **kwargs) -> CspNet:
     return _create_cspnet('darknet17', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def darknet21(pretrained=False, **kwargs):
+def darknet21(pretrained=False, **kwargs) -> CspNet:
     return _create_cspnet('darknet21', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def sedarknet21(pretrained=False, **kwargs):
+def sedarknet21(pretrained=False, **kwargs) -> CspNet:
     return _create_cspnet('sedarknet21', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def darknet53(pretrained=False, **kwargs):
+def darknet53(pretrained=False, **kwargs) -> CspNet:
     return _create_cspnet('darknet53', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def darknetaa53(pretrained=False, **kwargs):
+def darknetaa53(pretrained=False, **kwargs) -> CspNet:
     return _create_cspnet('darknetaa53', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def cs3darknet_s(pretrained=False, **kwargs):
+def cs3darknet_s(pretrained=False, **kwargs) -> CspNet:
     return _create_cspnet('cs3darknet_s', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def cs3darknet_m(pretrained=False, **kwargs):
+def cs3darknet_m(pretrained=False, **kwargs) -> CspNet:
     return _create_cspnet('cs3darknet_m', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def cs3darknet_l(pretrained=False, **kwargs):
+def cs3darknet_l(pretrained=False, **kwargs) -> CspNet:
     return _create_cspnet('cs3darknet_l', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def cs3darknet_x(pretrained=False, **kwargs):
+def cs3darknet_x(pretrained=False, **kwargs) -> CspNet:
     return _create_cspnet('cs3darknet_x', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def cs3darknet_focus_s(pretrained=False, **kwargs):
+def cs3darknet_focus_s(pretrained=False, **kwargs) -> CspNet:
     return _create_cspnet('cs3darknet_focus_s', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def cs3darknet_focus_m(pretrained=False, **kwargs):
+def cs3darknet_focus_m(pretrained=False, **kwargs) -> CspNet:
     return _create_cspnet('cs3darknet_focus_m', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def cs3darknet_focus_l(pretrained=False, **kwargs):
+def cs3darknet_focus_l(pretrained=False, **kwargs) -> CspNet:
     return _create_cspnet('cs3darknet_focus_l', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def cs3darknet_focus_x(pretrained=False, **kwargs):
+def cs3darknet_focus_x(pretrained=False, **kwargs) -> CspNet:
     return _create_cspnet('cs3darknet_focus_x', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def cs3sedarknet_l(pretrained=False, **kwargs):
+def cs3sedarknet_l(pretrained=False, **kwargs) -> CspNet:
     return _create_cspnet('cs3sedarknet_l', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def cs3sedarknet_x(pretrained=False, **kwargs):
+def cs3sedarknet_x(pretrained=False, **kwargs) -> CspNet:
     return _create_cspnet('cs3sedarknet_x', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def cs3sedarknet_xdw(pretrained=False, **kwargs):
+def cs3sedarknet_xdw(pretrained=False, **kwargs) -> CspNet:
     return _create_cspnet('cs3sedarknet_xdw', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def cs3edgenet_x(pretrained=False, **kwargs):
+def cs3edgenet_x(pretrained=False, **kwargs) -> CspNet:
     return _create_cspnet('cs3edgenet_x', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def cs3se_edgenet_x(pretrained=False, **kwargs):
+def cs3se_edgenet_x(pretrained=False, **kwargs) -> CspNet:
     return _create_cspnet('cs3se_edgenet_x', pretrained=pretrained, **kwargs)
\ No newline at end of file
diff --git a/timm/models/davit.py b/timm/models/davit.py
index 181d84b9f0..16d39b9d1b 100644
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@@ -27,7 +27,7 @@
 from ._manipulate import checkpoint_seq
 from ._registry import generate_default_cfgs, register_model
 
-__all__ = ['DaViT']
+__all__ = ['DaVit']
 
 
 class ConvPosEnc(nn.Module):
@@ -362,7 +362,7 @@ def forward(self, x: Tensor):
         return x
 
 
-class DaViTStage(nn.Module):
+class DaVitStage(nn.Module):
     def __init__(
             self,
             in_chs,
@@ -440,7 +440,7 @@ def forward(self, x: Tensor):
         return x
 
 
-class DaViT(nn.Module):
+class DaVit(nn.Module):
     r""" DaViT
         A PyTorch implementation of `DaViT: Dual Attention Vision Transformers`  - https://arxiv.org/abs/2204.03645
         Supports arbitrary input sizes and pyramid feature extraction
@@ -497,7 +497,7 @@ def __init__(
         stages = []
         for stage_idx in range(num_stages):
             out_chs = embed_dims[stage_idx]
-            stage = DaViTStage(
+            stage = DaVitStage(
                 in_chs,
                 out_chs,
                 depth=depths[stage_idx],
@@ -610,7 +610,7 @@ def _create_davit(variant, pretrained=False, **kwargs):
     out_indices = kwargs.pop('out_indices', default_out_indices)
 
     model = build_model_with_cfg(
-        DaViT,
+        DaVit,
         variant,
         pretrained,
         pretrained_filter_fn=checkpoint_filter_fn,
@@ -647,42 +647,42 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def davit_tiny(pretrained=False, **kwargs):
+def davit_tiny(pretrained=False, **kwargs) -> DaVit:
     model_kwargs = dict(
         depths=(1, 1, 3, 1), embed_dims=(96, 192, 384, 768), num_heads=(3, 6, 12, 24), **kwargs)
     return _create_davit('davit_tiny', pretrained=pretrained, **model_kwargs)
 
 
 @register_model
-def davit_small(pretrained=False, **kwargs):
+def davit_small(pretrained=False, **kwargs) -> DaVit:
     model_kwargs = dict(
         depths=(1, 1, 9, 1), embed_dims=(96, 192, 384, 768), num_heads=(3, 6, 12, 24), **kwargs)
     return _create_davit('davit_small', pretrained=pretrained, **model_kwargs)
 
 
 @register_model
-def davit_base(pretrained=False, **kwargs):
+def davit_base(pretrained=False, **kwargs) -> DaVit:
     model_kwargs = dict(
         depths=(1, 1, 9, 1), embed_dims=(128, 256, 512, 1024), num_heads=(4, 8, 16, 32), **kwargs)
     return _create_davit('davit_base', pretrained=pretrained, **model_kwargs)
 
 
 @register_model
-def davit_large(pretrained=False, **kwargs):
+def davit_large(pretrained=False, **kwargs) -> DaVit:
     model_kwargs = dict(
         depths=(1, 1, 9, 1), embed_dims=(192, 384, 768, 1536), num_heads=(6, 12, 24, 48), **kwargs)
     return _create_davit('davit_large', pretrained=pretrained, **model_kwargs)
 
 
 @register_model
-def davit_huge(pretrained=False, **kwargs):
+def davit_huge(pretrained=False, **kwargs) -> DaVit:
     model_kwargs = dict(
         depths=(1, 1, 9, 1), embed_dims=(256, 512, 1024, 2048), num_heads=(8, 16, 32, 64), **kwargs)
     return _create_davit('davit_huge', pretrained=pretrained, **model_kwargs)
 
 
 @register_model
-def davit_giant(pretrained=False, **kwargs):
+def davit_giant(pretrained=False, **kwargs) -> DaVit:
     model_kwargs = dict(
         depths=(1, 1, 12, 3), embed_dims=(384, 768, 1536, 3072), num_heads=(12, 24, 48, 96), **kwargs)
     return _create_davit('davit_giant', pretrained=pretrained, **model_kwargs)
diff --git a/timm/models/deit.py b/timm/models/deit.py
index 8c3bd38d9c..be52a977e6 100644
--- a/timm/models/deit.py
+++ b/timm/models/deit.py
@@ -222,7 +222,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def deit_tiny_patch16_224(pretrained=False, **kwargs):
+def deit_tiny_patch16_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ DeiT-tiny model @ 224x224 from paper (https://arxiv.org/abs/2012.12877).
     ImageNet-1k weights from https://github.com/facebookresearch/deit.
     """
@@ -232,7 +232,7 @@ def deit_tiny_patch16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def deit_small_patch16_224(pretrained=False, **kwargs):
+def deit_small_patch16_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ DeiT-small model @ 224x224 from paper (https://arxiv.org/abs/2012.12877).
     ImageNet-1k weights from https://github.com/facebookresearch/deit.
     """
@@ -242,7 +242,7 @@ def deit_small_patch16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def deit_base_patch16_224(pretrained=False, **kwargs):
+def deit_base_patch16_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ DeiT base model @ 224x224 from paper (https://arxiv.org/abs/2012.12877).
     ImageNet-1k weights from https://github.com/facebookresearch/deit.
     """
@@ -252,7 +252,7 @@ def deit_base_patch16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def deit_base_patch16_384(pretrained=False, **kwargs):
+def deit_base_patch16_384(pretrained=False, **kwargs) -> VisionTransformer:
     """ DeiT base model @ 384x384 from paper (https://arxiv.org/abs/2012.12877).
     ImageNet-1k weights from https://github.com/facebookresearch/deit.
     """
@@ -262,7 +262,7 @@ def deit_base_patch16_384(pretrained=False, **kwargs):
 
 
 @register_model
-def deit_tiny_distilled_patch16_224(pretrained=False, **kwargs):
+def deit_tiny_distilled_patch16_224(pretrained=False, **kwargs) -> VisionTransformerDistilled:
     """ DeiT-tiny distilled model @ 224x224 from paper (https://arxiv.org/abs/2012.12877).
     ImageNet-1k weights from https://github.com/facebookresearch/deit.
     """
@@ -273,7 +273,7 @@ def deit_tiny_distilled_patch16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def deit_small_distilled_patch16_224(pretrained=False, **kwargs):
+def deit_small_distilled_patch16_224(pretrained=False, **kwargs) -> VisionTransformerDistilled:
     """ DeiT-small distilled model @ 224x224 from paper (https://arxiv.org/abs/2012.12877).
     ImageNet-1k weights from https://github.com/facebookresearch/deit.
     """
@@ -284,7 +284,7 @@ def deit_small_distilled_patch16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def deit_base_distilled_patch16_224(pretrained=False, **kwargs):
+def deit_base_distilled_patch16_224(pretrained=False, **kwargs) -> VisionTransformerDistilled:
     """ DeiT-base distilled model @ 224x224 from paper (https://arxiv.org/abs/2012.12877).
     ImageNet-1k weights from https://github.com/facebookresearch/deit.
     """
@@ -295,7 +295,7 @@ def deit_base_distilled_patch16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def deit_base_distilled_patch16_384(pretrained=False, **kwargs):
+def deit_base_distilled_patch16_384(pretrained=False, **kwargs) -> VisionTransformerDistilled:
     """ DeiT-base distilled model @ 384x384 from paper (https://arxiv.org/abs/2012.12877).
     ImageNet-1k weights from https://github.com/facebookresearch/deit.
     """
@@ -306,7 +306,7 @@ def deit_base_distilled_patch16_384(pretrained=False, **kwargs):
 
 
 @register_model
-def deit3_small_patch16_224(pretrained=False, **kwargs):
+def deit3_small_patch16_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ DeiT-3 small model @ 224x224 from paper (https://arxiv.org/abs/2204.07118).
     ImageNet-1k weights from https://github.com/facebookresearch/deit.
     """
@@ -316,7 +316,7 @@ def deit3_small_patch16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def deit3_small_patch16_384(pretrained=False, **kwargs):
+def deit3_small_patch16_384(pretrained=False, **kwargs) -> VisionTransformer:
     """ DeiT-3 small model @ 384x384 from paper (https://arxiv.org/abs/2204.07118).
     ImageNet-1k weights from https://github.com/facebookresearch/deit.
     """
@@ -326,7 +326,7 @@ def deit3_small_patch16_384(pretrained=False, **kwargs):
 
 
 @register_model
-def deit3_medium_patch16_224(pretrained=False, **kwargs):
+def deit3_medium_patch16_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ DeiT-3 medium model @ 224x224 (https://arxiv.org/abs/2012.12877).
     ImageNet-1k weights from https://github.com/facebookresearch/deit.
     """
@@ -336,7 +336,7 @@ def deit3_medium_patch16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def deit3_base_patch16_224(pretrained=False, **kwargs):
+def deit3_base_patch16_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ DeiT-3 base model @ 224x224 from paper (https://arxiv.org/abs/2204.07118).
     ImageNet-1k weights from https://github.com/facebookresearch/deit.
     """
@@ -346,7 +346,7 @@ def deit3_base_patch16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def deit3_base_patch16_384(pretrained=False, **kwargs):
+def deit3_base_patch16_384(pretrained=False, **kwargs) -> VisionTransformer:
     """ DeiT-3 base model @ 384x384 from paper (https://arxiv.org/abs/2204.07118).
     ImageNet-1k weights from https://github.com/facebookresearch/deit.
     """
@@ -356,7 +356,7 @@ def deit3_base_patch16_384(pretrained=False, **kwargs):
 
 
 @register_model
-def deit3_large_patch16_224(pretrained=False, **kwargs):
+def deit3_large_patch16_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ DeiT-3 large model @ 224x224 from paper (https://arxiv.org/abs/2204.07118).
     ImageNet-1k weights from https://github.com/facebookresearch/deit.
     """
@@ -366,7 +366,7 @@ def deit3_large_patch16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def deit3_large_patch16_384(pretrained=False, **kwargs):
+def deit3_large_patch16_384(pretrained=False, **kwargs) -> VisionTransformer:
     """ DeiT-3 large model @ 384x384 from paper (https://arxiv.org/abs/2204.07118).
     ImageNet-1k weights from https://github.com/facebookresearch/deit.
     """
@@ -376,7 +376,7 @@ def deit3_large_patch16_384(pretrained=False, **kwargs):
 
 
 @register_model
-def deit3_huge_patch14_224(pretrained=False, **kwargs):
+def deit3_huge_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ DeiT-3 base model @ 384x384 from paper (https://arxiv.org/abs/2204.07118).
     ImageNet-1k weights from https://github.com/facebookresearch/deit.
     """
diff --git a/timm/models/densenet.py b/timm/models/densenet.py
index 40beeb703d..ab34d6930f 100644
--- a/timm/models/densenet.py
+++ b/timm/models/densenet.py
@@ -357,7 +357,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def densenet121(pretrained=False, **kwargs):
+def densenet121(pretrained=False, **kwargs) -> DenseNet:
     r"""Densenet-121 model from
     `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`
     """
@@ -367,7 +367,7 @@ def densenet121(pretrained=False, **kwargs):
 
 
 @register_model
-def densenetblur121d(pretrained=False, **kwargs):
+def densenetblur121d(pretrained=False, **kwargs) -> DenseNet:
     r"""Densenet-121 w/ blur-pooling & 3-layer 3x3 stem
     `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`
     """
@@ -378,7 +378,7 @@ def densenetblur121d(pretrained=False, **kwargs):
 
 
 @register_model
-def densenet169(pretrained=False, **kwargs):
+def densenet169(pretrained=False, **kwargs) -> DenseNet:
     r"""Densenet-169 model from
     `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`
     """
@@ -388,7 +388,7 @@ def densenet169(pretrained=False, **kwargs):
 
 
 @register_model
-def densenet201(pretrained=False, **kwargs):
+def densenet201(pretrained=False, **kwargs) -> DenseNet:
     r"""Densenet-201 model from
     `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`
     """
@@ -398,7 +398,7 @@ def densenet201(pretrained=False, **kwargs):
 
 
 @register_model
-def densenet161(pretrained=False, **kwargs):
+def densenet161(pretrained=False, **kwargs) -> DenseNet:
     r"""Densenet-161 model from
     `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`
     """
@@ -408,7 +408,7 @@ def densenet161(pretrained=False, **kwargs):
 
 
 @register_model
-def densenet264d(pretrained=False, **kwargs):
+def densenet264d(pretrained=False, **kwargs) -> DenseNet:
     r"""Densenet-264 model from
     `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`
     """
diff --git a/timm/models/dla.py b/timm/models/dla.py
index 565798d974..e7c20dca03 100644
--- a/timm/models/dla.py
+++ b/timm/models/dla.py
@@ -422,7 +422,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def dla60_res2net(pretrained=False, **kwargs):
+def dla60_res2net(pretrained=False, **kwargs) -> DLA:
     model_args = dict(
         levels=(1, 1, 1, 2, 3, 1), channels=(16, 32, 128, 256, 512, 1024),
         block=DlaBottle2neck, cardinality=1, base_width=28)
@@ -438,21 +438,21 @@ def dla60_res2next(pretrained=False,**kwargs):
 
 
 @register_model
-def dla34(pretrained=False, **kwargs):  # DLA-34
+def dla34(pretrained=False, **kwargs) -> DLA:  # DLA-34
     model_args = dict(
         levels=[1, 1, 1, 2, 2, 1], channels=[16, 32, 64, 128, 256, 512], block=DlaBasic)
     return _create_dla('dla34', pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
-def dla46_c(pretrained=False, **kwargs):  # DLA-46-C
+def dla46_c(pretrained=False, **kwargs) -> DLA:  # DLA-46-C
     model_args = dict(
         levels=[1, 1, 1, 2, 2, 1], channels=[16, 32, 64, 64, 128, 256], block=DlaBottleneck)
     return _create_dla('dla46_c', pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
-def dla46x_c(pretrained=False, **kwargs):  # DLA-X-46-C
+def dla46x_c(pretrained=False, **kwargs) -> DLA:  # DLA-X-46-C
     model_args = dict(
         levels=[1, 1, 1, 2, 2, 1], channels=[16, 32, 64, 64, 128, 256],
         block=DlaBottleneck, cardinality=32, base_width=4)
@@ -460,7 +460,7 @@ def dla46x_c(pretrained=False, **kwargs):  # DLA-X-46-C
 
 
 @register_model
-def dla60x_c(pretrained=False, **kwargs):  # DLA-X-60-C
+def dla60x_c(pretrained=False, **kwargs) -> DLA:  # DLA-X-60-C
     model_args = dict(
         levels=[1, 1, 1, 2, 3, 1], channels=[16, 32, 64, 64, 128, 256],
         block=DlaBottleneck, cardinality=32, base_width=4)
@@ -468,7 +468,7 @@ def dla60x_c(pretrained=False, **kwargs):  # DLA-X-60-C
 
 
 @register_model
-def dla60(pretrained=False, **kwargs):  # DLA-60
+def dla60(pretrained=False, **kwargs) -> DLA:  # DLA-60
     model_args = dict(
         levels=[1, 1, 1, 2, 3, 1], channels=[16, 32, 128, 256, 512, 1024],
         block=DlaBottleneck)
@@ -476,7 +476,7 @@ def dla60(pretrained=False, **kwargs):  # DLA-60
 
 
 @register_model
-def dla60x(pretrained=False, **kwargs):  # DLA-X-60
+def dla60x(pretrained=False, **kwargs) -> DLA:  # DLA-X-60
     model_args = dict(
         levels=[1, 1, 1, 2, 3, 1], channels=[16, 32, 128, 256, 512, 1024],
         block=DlaBottleneck, cardinality=32, base_width=4)
@@ -484,7 +484,7 @@ def dla60x(pretrained=False, **kwargs):  # DLA-X-60
 
 
 @register_model
-def dla102(pretrained=False, **kwargs):  # DLA-102
+def dla102(pretrained=False, **kwargs) -> DLA:  # DLA-102
     model_args = dict(
         levels=[1, 1, 1, 3, 4, 1], channels=[16, 32, 128, 256, 512, 1024],
         block=DlaBottleneck, shortcut_root=True)
@@ -492,7 +492,7 @@ def dla102(pretrained=False, **kwargs):  # DLA-102
 
 
 @register_model
-def dla102x(pretrained=False, **kwargs):  # DLA-X-102
+def dla102x(pretrained=False, **kwargs) -> DLA:  # DLA-X-102
     model_args = dict(
         levels=[1, 1, 1, 3, 4, 1], channels=[16, 32, 128, 256, 512, 1024],
         block=DlaBottleneck, cardinality=32, base_width=4, shortcut_root=True)
@@ -500,7 +500,7 @@ def dla102x(pretrained=False, **kwargs):  # DLA-X-102
 
 
 @register_model
-def dla102x2(pretrained=False, **kwargs):  # DLA-X-102 64
+def dla102x2(pretrained=False, **kwargs) -> DLA:  # DLA-X-102 64
     model_args = dict(
         levels=[1, 1, 1, 3, 4, 1], channels=[16, 32, 128, 256, 512, 1024],
         block=DlaBottleneck, cardinality=64, base_width=4, shortcut_root=True)
@@ -508,7 +508,7 @@ def dla102x2(pretrained=False, **kwargs):  # DLA-X-102 64
 
 
 @register_model
-def dla169(pretrained=False, **kwargs):  # DLA-169
+def dla169(pretrained=False, **kwargs) -> DLA:  # DLA-169
     model_args = dict(
         levels=[1, 1, 2, 3, 5, 1], channels=[16, 32, 128, 256, 512, 1024],
         block=DlaBottleneck, shortcut_root=True)
diff --git a/timm/models/dpn.py b/timm/models/dpn.py
index 31ae00a56c..a7ea517e0d 100644
--- a/timm/models/dpn.py
+++ b/timm/models/dpn.py
@@ -316,7 +316,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def dpn48b(pretrained=False, **kwargs):
+def dpn48b(pretrained=False, **kwargs) -> DPN:
     model_kwargs = dict(
         small=True, num_init_features=10, k_r=128, groups=32,
         b=True, k_sec=(3, 4, 6, 3), inc_sec=(16, 32, 32, 64), act_layer='silu')
@@ -324,7 +324,7 @@ def dpn48b(pretrained=False, **kwargs):
 
 
 @register_model
-def dpn68(pretrained=False, **kwargs):
+def dpn68(pretrained=False, **kwargs) -> DPN:
     model_kwargs = dict(
         small=True, num_init_features=10, k_r=128, groups=32,
         k_sec=(3, 4, 12, 3), inc_sec=(16, 32, 32, 64))
@@ -332,7 +332,7 @@ def dpn68(pretrained=False, **kwargs):
 
 
 @register_model
-def dpn68b(pretrained=False, **kwargs):
+def dpn68b(pretrained=False, **kwargs) -> DPN:
     model_kwargs = dict(
         small=True, num_init_features=10, k_r=128, groups=32,
         b=True, k_sec=(3, 4, 12, 3), inc_sec=(16, 32, 32, 64))
@@ -340,7 +340,7 @@ def dpn68b(pretrained=False, **kwargs):
 
 
 @register_model
-def dpn92(pretrained=False, **kwargs):
+def dpn92(pretrained=False, **kwargs) -> DPN:
     model_kwargs = dict(
         num_init_features=64, k_r=96, groups=32,
         k_sec=(3, 4, 20, 3), inc_sec=(16, 32, 24, 128))
@@ -348,7 +348,7 @@ def dpn92(pretrained=False, **kwargs):
 
 
 @register_model
-def dpn98(pretrained=False, **kwargs):
+def dpn98(pretrained=False, **kwargs) -> DPN:
     model_kwargs = dict(
         num_init_features=96, k_r=160, groups=40,
         k_sec=(3, 6, 20, 3), inc_sec=(16, 32, 32, 128))
@@ -356,7 +356,7 @@ def dpn98(pretrained=False, **kwargs):
 
 
 @register_model
-def dpn131(pretrained=False, **kwargs):
+def dpn131(pretrained=False, **kwargs) -> DPN:
     model_kwargs = dict(
         num_init_features=128, k_r=160, groups=40,
         k_sec=(4, 8, 28, 3), inc_sec=(16, 32, 32, 128))
@@ -364,7 +364,7 @@ def dpn131(pretrained=False, **kwargs):
 
 
 @register_model
-def dpn107(pretrained=False, **kwargs):
+def dpn107(pretrained=False, **kwargs) -> DPN:
     model_kwargs = dict(
         num_init_features=128, k_r=200, groups=50,
         k_sec=(4, 8, 20, 3), inc_sec=(20, 64, 64, 128))
diff --git a/timm/models/edgenext.py b/timm/models/edgenext.py
index c597b7ed12..b6b15c1229 100644
--- a/timm/models/edgenext.py
+++ b/timm/models/edgenext.py
@@ -522,7 +522,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def edgenext_xx_small(pretrained=False, **kwargs):
+def edgenext_xx_small(pretrained=False, **kwargs) -> EdgeNeXt:
     # 1.33M & 260.58M @ 256 resolution
     # 71.23% Top-1 accuracy
     # No AA, Color Jitter=0.4, No Mixup & Cutmix, DropPath=0.0, BS=4096, lr=0.006, multi-scale-sampler
@@ -533,7 +533,7 @@ def edgenext_xx_small(pretrained=False, **kwargs):
 
 
 @register_model
-def edgenext_x_small(pretrained=False, **kwargs):
+def edgenext_x_small(pretrained=False, **kwargs) -> EdgeNeXt:
     # 2.34M & 538.0M @ 256 resolution
     # 75.00% Top-1 accuracy
     # No AA, No Mixup & Cutmix, DropPath=0.0, BS=4096, lr=0.006, multi-scale-sampler
@@ -544,7 +544,7 @@ def edgenext_x_small(pretrained=False, **kwargs):
 
 
 @register_model
-def edgenext_small(pretrained=False, **kwargs):
+def edgenext_small(pretrained=False, **kwargs) -> EdgeNeXt:
     # 5.59M & 1260.59M @ 256 resolution
     # 79.43% Top-1 accuracy
     # AA=True, No Mixup & Cutmix, DropPath=0.1, BS=4096, lr=0.006, multi-scale-sampler
@@ -555,7 +555,7 @@ def edgenext_small(pretrained=False, **kwargs):
 
 
 @register_model
-def edgenext_base(pretrained=False, **kwargs):
+def edgenext_base(pretrained=False, **kwargs) -> EdgeNeXt:
     # 18.51M & 3840.93M @ 256 resolution
     # 82.5% (normal) 83.7% (USI) Top-1 accuracy
     # AA=True, Mixup & Cutmix, DropPath=0.1, BS=4096, lr=0.006, multi-scale-sampler
@@ -566,7 +566,7 @@ def edgenext_base(pretrained=False, **kwargs):
 
 
 @register_model
-def edgenext_small_rw(pretrained=False, **kwargs):
+def edgenext_small_rw(pretrained=False, **kwargs) -> EdgeNeXt:
     model_kwargs = dict(
         depths=(3, 3, 9, 3), dims=(48, 96, 192, 384),
         downsample_block=True, conv_bias=False, stem_type='overlap', **kwargs)
diff --git a/timm/models/efficientformer.py b/timm/models/efficientformer.py
index 62ff36de2a..970d1a112e 100644
--- a/timm/models/efficientformer.py
+++ b/timm/models/efficientformer.py
@@ -542,7 +542,7 @@ def _create_efficientformer(variant, pretrained=False, **kwargs):
 
 
 @register_model
-def efficientformer_l1(pretrained=False, **kwargs):
+def efficientformer_l1(pretrained=False, **kwargs) -> EfficientFormer:
     model_args = dict(
         depths=EfficientFormer_depth['l1'],
         embed_dims=EfficientFormer_width['l1'],
@@ -552,7 +552,7 @@ def efficientformer_l1(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientformer_l3(pretrained=False, **kwargs):
+def efficientformer_l3(pretrained=False, **kwargs) -> EfficientFormer:
     model_args = dict(
         depths=EfficientFormer_depth['l3'],
         embed_dims=EfficientFormer_width['l3'],
@@ -562,7 +562,7 @@ def efficientformer_l3(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientformer_l7(pretrained=False, **kwargs):
+def efficientformer_l7(pretrained=False, **kwargs) -> EfficientFormer:
     model_args = dict(
         depths=EfficientFormer_depth['l7'],
         embed_dims=EfficientFormer_width['l7'],
diff --git a/timm/models/efficientformer_v2.py b/timm/models/efficientformer_v2.py
index 166aaaaee8..9388131ed6 100644
--- a/timm/models/efficientformer_v2.py
+++ b/timm/models/efficientformer_v2.py
@@ -687,7 +687,7 @@ def _create_efficientformerv2(variant, pretrained=False, **kwargs):
 
 
 @register_model
-def efficientformerv2_s0(pretrained=False, **kwargs):
+def efficientformerv2_s0(pretrained=False, **kwargs) -> EfficientFormerV2:
     model_args = dict(
         depths=EfficientFormer_depth['S0'],
         embed_dims=EfficientFormer_width['S0'],
@@ -699,7 +699,7 @@ def efficientformerv2_s0(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientformerv2_s1(pretrained=False, **kwargs):
+def efficientformerv2_s1(pretrained=False, **kwargs) -> EfficientFormerV2:
     model_args = dict(
         depths=EfficientFormer_depth['S1'],
         embed_dims=EfficientFormer_width['S1'],
@@ -711,7 +711,7 @@ def efficientformerv2_s1(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientformerv2_s2(pretrained=False, **kwargs):
+def efficientformerv2_s2(pretrained=False, **kwargs) -> EfficientFormerV2:
     model_args = dict(
         depths=EfficientFormer_depth['S2'],
         embed_dims=EfficientFormer_width['S2'],
@@ -723,7 +723,7 @@ def efficientformerv2_s2(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientformerv2_l(pretrained=False, **kwargs):
+def efficientformerv2_l(pretrained=False, **kwargs) -> EfficientFormerV2:
     model_args = dict(
         depths=EfficientFormer_depth['L'],
         embed_dims=EfficientFormer_width['L'],
diff --git a/timm/models/efficientnet.py b/timm/models/efficientnet.py
index 3a9fc13a6a..1854bc3058 100644
--- a/timm/models/efficientnet.py
+++ b/timm/models/efficientnet.py
@@ -1413,117 +1413,117 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def mnasnet_050(pretrained=False, **kwargs):
+def mnasnet_050(pretrained=False, **kwargs) -> EfficientNet:
     """ MNASNet B1, depth multiplier of 0.5. """
     model = _gen_mnasnet_b1('mnasnet_050', 0.5, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def mnasnet_075(pretrained=False, **kwargs):
+def mnasnet_075(pretrained=False, **kwargs) -> EfficientNet:
     """ MNASNet B1, depth multiplier of 0.75. """
     model = _gen_mnasnet_b1('mnasnet_075', 0.75, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def mnasnet_100(pretrained=False, **kwargs):
+def mnasnet_100(pretrained=False, **kwargs) -> EfficientNet:
     """ MNASNet B1, depth multiplier of 1.0. """
     model = _gen_mnasnet_b1('mnasnet_100', 1.0, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def mnasnet_b1(pretrained=False, **kwargs):
+def mnasnet_b1(pretrained=False, **kwargs) -> EfficientNet:
     """ MNASNet B1, depth multiplier of 1.0. """
     return mnasnet_100(pretrained, **kwargs)
 
 
 @register_model
-def mnasnet_140(pretrained=False, **kwargs):
+def mnasnet_140(pretrained=False, **kwargs) -> EfficientNet:
     """ MNASNet B1,  depth multiplier of 1.4 """
     model = _gen_mnasnet_b1('mnasnet_140', 1.4, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def semnasnet_050(pretrained=False, **kwargs):
+def semnasnet_050(pretrained=False, **kwargs) -> EfficientNet:
     """ MNASNet A1 (w/ SE), depth multiplier of 0.5 """
     model = _gen_mnasnet_a1('semnasnet_050', 0.5, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def semnasnet_075(pretrained=False, **kwargs):
+def semnasnet_075(pretrained=False, **kwargs) -> EfficientNet:
     """ MNASNet A1 (w/ SE),  depth multiplier of 0.75. """
     model = _gen_mnasnet_a1('semnasnet_075', 0.75, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def semnasnet_100(pretrained=False, **kwargs):
+def semnasnet_100(pretrained=False, **kwargs) -> EfficientNet:
     """ MNASNet A1 (w/ SE), depth multiplier of 1.0. """
     model = _gen_mnasnet_a1('semnasnet_100', 1.0, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def mnasnet_a1(pretrained=False, **kwargs):
+def mnasnet_a1(pretrained=False, **kwargs) -> EfficientNet:
     """ MNASNet A1 (w/ SE), depth multiplier of 1.0. """
     return semnasnet_100(pretrained, **kwargs)
 
 
 @register_model
-def semnasnet_140(pretrained=False, **kwargs):
+def semnasnet_140(pretrained=False, **kwargs) -> EfficientNet:
     """ MNASNet A1 (w/ SE), depth multiplier of 1.4. """
     model = _gen_mnasnet_a1('semnasnet_140', 1.4, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def mnasnet_small(pretrained=False, **kwargs):
+def mnasnet_small(pretrained=False, **kwargs) -> EfficientNet:
     """ MNASNet Small,  depth multiplier of 1.0. """
     model = _gen_mnasnet_small('mnasnet_small', 1.0, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def mobilenetv2_035(pretrained=False, **kwargs):
+def mobilenetv2_035(pretrained=False, **kwargs) -> EfficientNet:
     """ MobileNet V2 w/ 0.35 channel multiplier """
     model = _gen_mobilenet_v2('mobilenetv2_035', 0.35, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def mobilenetv2_050(pretrained=False, **kwargs):
+def mobilenetv2_050(pretrained=False, **kwargs) -> EfficientNet:
     """ MobileNet V2 w/ 0.5 channel multiplier """
     model = _gen_mobilenet_v2('mobilenetv2_050', 0.5, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def mobilenetv2_075(pretrained=False, **kwargs):
+def mobilenetv2_075(pretrained=False, **kwargs) -> EfficientNet:
     """ MobileNet V2 w/ 0.75 channel multiplier """
     model = _gen_mobilenet_v2('mobilenetv2_075', 0.75, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def mobilenetv2_100(pretrained=False, **kwargs):
+def mobilenetv2_100(pretrained=False, **kwargs) -> EfficientNet:
     """ MobileNet V2 w/ 1.0 channel multiplier """
     model = _gen_mobilenet_v2('mobilenetv2_100', 1.0, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def mobilenetv2_140(pretrained=False, **kwargs):
+def mobilenetv2_140(pretrained=False, **kwargs) -> EfficientNet:
     """ MobileNet V2 w/ 1.4 channel multiplier """
     model = _gen_mobilenet_v2('mobilenetv2_140', 1.4, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def mobilenetv2_110d(pretrained=False, **kwargs):
+def mobilenetv2_110d(pretrained=False, **kwargs) -> EfficientNet:
     """ MobileNet V2 w/ 1.1 channel, 1.2 depth multipliers"""
     model = _gen_mobilenet_v2(
         'mobilenetv2_110d', 1.1, depth_multiplier=1.2, fix_stem_head=True, pretrained=pretrained, **kwargs)
@@ -1531,7 +1531,7 @@ def mobilenetv2_110d(pretrained=False, **kwargs):
 
 
 @register_model
-def mobilenetv2_120d(pretrained=False, **kwargs):
+def mobilenetv2_120d(pretrained=False, **kwargs) -> EfficientNet:
     """ MobileNet V2 w/ 1.2 channel, 1.4 depth multipliers """
     model = _gen_mobilenet_v2(
         'mobilenetv2_120d', 1.2, depth_multiplier=1.4, fix_stem_head=True, pretrained=pretrained, **kwargs)
@@ -1539,7 +1539,7 @@ def mobilenetv2_120d(pretrained=False, **kwargs):
 
 
 @register_model
-def fbnetc_100(pretrained=False, **kwargs):
+def fbnetc_100(pretrained=False, **kwargs) -> EfficientNet:
     """ FBNet-C """
     if pretrained:
         # pretrained model trained with non-default BN epsilon
@@ -1549,14 +1549,14 @@ def fbnetc_100(pretrained=False, **kwargs):
 
 
 @register_model
-def spnasnet_100(pretrained=False, **kwargs):
+def spnasnet_100(pretrained=False, **kwargs) -> EfficientNet:
     """ Single-Path NAS Pixel1"""
     model = _gen_spnasnet('spnasnet_100', 1.0, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def efficientnet_b0(pretrained=False, **kwargs):
+def efficientnet_b0(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B0 """
     # NOTE for train, drop_rate should be 0.2, drop_path_rate should be 0.2
     model = _gen_efficientnet(
@@ -1565,7 +1565,7 @@ def efficientnet_b0(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_b1(pretrained=False, **kwargs):
+def efficientnet_b1(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B1 """
     # NOTE for train, drop_rate should be 0.2, drop_path_rate should be 0.2
     model = _gen_efficientnet(
@@ -1574,7 +1574,7 @@ def efficientnet_b1(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_b2(pretrained=False, **kwargs):
+def efficientnet_b2(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B2 """
     # NOTE for train, drop_rate should be 0.3, drop_path_rate should be 0.2
     model = _gen_efficientnet(
@@ -1583,14 +1583,14 @@ def efficientnet_b2(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_b2a(pretrained=False, **kwargs):
+def efficientnet_b2a(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B2 @ 288x288 w/ 1.0 test crop"""
     # WARN this model def is deprecated, different train/test res + test crop handled by default_cfg now
     return efficientnet_b2(pretrained=pretrained, **kwargs)
 
 
 @register_model
-def efficientnet_b3(pretrained=False, **kwargs):
+def efficientnet_b3(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B3 """
     # NOTE for train, drop_rate should be 0.3, drop_path_rate should be 0.2
     model = _gen_efficientnet(
@@ -1599,14 +1599,14 @@ def efficientnet_b3(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_b3a(pretrained=False, **kwargs):
+def efficientnet_b3a(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B3 @ 320x320 w/ 1.0 test crop-pct """
     # WARN this model def is deprecated, different train/test res + test crop handled by default_cfg now
     return efficientnet_b3(pretrained=pretrained, **kwargs)
 
 
 @register_model
-def efficientnet_b4(pretrained=False, **kwargs):
+def efficientnet_b4(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B4 """
     # NOTE for train, drop_rate should be 0.4, drop_path_rate should be 0.2
     model = _gen_efficientnet(
@@ -1615,7 +1615,7 @@ def efficientnet_b4(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_b5(pretrained=False, **kwargs):
+def efficientnet_b5(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B5 """
     # NOTE for train, drop_rate should be 0.4, drop_path_rate should be 0.2
     model = _gen_efficientnet(
@@ -1624,7 +1624,7 @@ def efficientnet_b5(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_b6(pretrained=False, **kwargs):
+def efficientnet_b6(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B6 """
     # NOTE for train, drop_rate should be 0.5, drop_path_rate should be 0.2
     model = _gen_efficientnet(
@@ -1633,7 +1633,7 @@ def efficientnet_b6(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_b7(pretrained=False, **kwargs):
+def efficientnet_b7(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B7 """
     # NOTE for train, drop_rate should be 0.5, drop_path_rate should be 0.2
     model = _gen_efficientnet(
@@ -1642,7 +1642,7 @@ def efficientnet_b7(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_b8(pretrained=False, **kwargs):
+def efficientnet_b8(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B8 """
     # NOTE for train, drop_rate should be 0.5, drop_path_rate should be 0.2
     model = _gen_efficientnet(
@@ -1651,7 +1651,7 @@ def efficientnet_b8(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_l2(pretrained=False, **kwargs):
+def efficientnet_l2(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-L2."""
     # NOTE for train, drop_rate should be 0.5, drop_path_rate should be 0.2
     model = _gen_efficientnet(
@@ -1661,7 +1661,7 @@ def efficientnet_l2(pretrained=False, **kwargs):
 
 # FIXME experimental group cong / GroupNorm / EvoNorm experiments
 @register_model
-def efficientnet_b0_gn(pretrained=False, **kwargs):
+def efficientnet_b0_gn(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B0 + GroupNorm"""
     model = _gen_efficientnet(
         'efficientnet_b0_gn', norm_layer=partial(GroupNormAct, group_size=8), pretrained=pretrained, **kwargs)
@@ -1669,7 +1669,7 @@ def efficientnet_b0_gn(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_b0_g8_gn(pretrained=False, **kwargs):
+def efficientnet_b0_g8_gn(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B0 w/ group conv + GroupNorm"""
     model = _gen_efficientnet(
         'efficientnet_b0_g8_gn', group_size=8, norm_layer=partial(GroupNormAct, group_size=8),
@@ -1678,7 +1678,7 @@ def efficientnet_b0_g8_gn(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_b0_g16_evos(pretrained=False, **kwargs):
+def efficientnet_b0_g16_evos(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B0 w/ group 16 conv + EvoNorm"""
     model = _gen_efficientnet(
         'efficientnet_b0_g16_evos', group_size=16, channel_divisor=16,
@@ -1687,7 +1687,7 @@ def efficientnet_b0_g16_evos(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_b3_gn(pretrained=False, **kwargs):
+def efficientnet_b3_gn(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B3 w/ GroupNorm """
     # NOTE for train, drop_rate should be 0.3, drop_path_rate should be 0.2
     model = _gen_efficientnet(
@@ -1697,7 +1697,7 @@ def efficientnet_b3_gn(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_b3_g8_gn(pretrained=False, **kwargs):
+def efficientnet_b3_g8_gn(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B3 w/ grouped conv + BN"""
     # NOTE for train, drop_rate should be 0.3, drop_path_rate should be 0.2
     model = _gen_efficientnet(
@@ -1707,7 +1707,7 @@ def efficientnet_b3_g8_gn(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_es(pretrained=False, **kwargs):
+def efficientnet_es(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-Edge Small. """
     model = _gen_efficientnet_edge(
         'efficientnet_es', channel_multiplier=1.0, depth_multiplier=1.0, pretrained=pretrained, **kwargs)
@@ -1715,14 +1715,14 @@ def efficientnet_es(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_es_pruned(pretrained=False, **kwargs):
+def efficientnet_es_pruned(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-Edge Small Pruned. For more info: https://github.com/DeGirum/pruned-models/releases/tag/efficientnet_v1.0"""
     model = _gen_efficientnet_edge(
         'efficientnet_es_pruned', channel_multiplier=1.0, depth_multiplier=1.0, pretrained=pretrained, **kwargs)
     return model
 
 @register_model
-def efficientnet_em(pretrained=False, **kwargs):
+def efficientnet_em(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-Edge-Medium. """
     model = _gen_efficientnet_edge(
         'efficientnet_em', channel_multiplier=1.0, depth_multiplier=1.1, pretrained=pretrained, **kwargs)
@@ -1730,21 +1730,21 @@ def efficientnet_em(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_el(pretrained=False, **kwargs):
+def efficientnet_el(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-Edge-Large. """
     model = _gen_efficientnet_edge(
         'efficientnet_el', channel_multiplier=1.2, depth_multiplier=1.4, pretrained=pretrained, **kwargs)
     return model
 
 @register_model
-def efficientnet_el_pruned(pretrained=False, **kwargs):
+def efficientnet_el_pruned(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-Edge-Large pruned. For more info: https://github.com/DeGirum/pruned-models/releases/tag/efficientnet_v1.0"""
     model = _gen_efficientnet_edge(
         'efficientnet_el_pruned', channel_multiplier=1.2, depth_multiplier=1.4, pretrained=pretrained, **kwargs)
     return model
 
 @register_model
-def efficientnet_cc_b0_4e(pretrained=False, **kwargs):
+def efficientnet_cc_b0_4e(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-CondConv-B0 w/ 8 Experts """
     # NOTE for train, drop_rate should be 0.2, drop_path_rate should be 0.2
     model = _gen_efficientnet_condconv(
@@ -1753,7 +1753,7 @@ def efficientnet_cc_b0_4e(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_cc_b0_8e(pretrained=False, **kwargs):
+def efficientnet_cc_b0_8e(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-CondConv-B0 w/ 8 Experts """
     # NOTE for train, drop_rate should be 0.2, drop_path_rate should be 0.2
     model = _gen_efficientnet_condconv(
@@ -1763,7 +1763,7 @@ def efficientnet_cc_b0_8e(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_cc_b1_8e(pretrained=False, **kwargs):
+def efficientnet_cc_b1_8e(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-CondConv-B1 w/ 8 Experts """
     # NOTE for train, drop_rate should be 0.2, drop_path_rate should be 0.2
     model = _gen_efficientnet_condconv(
@@ -1773,7 +1773,7 @@ def efficientnet_cc_b1_8e(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_lite0(pretrained=False, **kwargs):
+def efficientnet_lite0(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-Lite0 """
     # NOTE for train, drop_rate should be 0.2, drop_path_rate should be 0.2
     model = _gen_efficientnet_lite(
@@ -1782,7 +1782,7 @@ def efficientnet_lite0(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_lite1(pretrained=False, **kwargs):
+def efficientnet_lite1(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-Lite1 """
     # NOTE for train, drop_rate should be 0.2, drop_path_rate should be 0.2
     model = _gen_efficientnet_lite(
@@ -1791,7 +1791,7 @@ def efficientnet_lite1(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_lite2(pretrained=False, **kwargs):
+def efficientnet_lite2(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-Lite2 """
     # NOTE for train, drop_rate should be 0.3, drop_path_rate should be 0.2
     model = _gen_efficientnet_lite(
@@ -1800,7 +1800,7 @@ def efficientnet_lite2(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_lite3(pretrained=False, **kwargs):
+def efficientnet_lite3(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-Lite3 """
     # NOTE for train, drop_rate should be 0.3, drop_path_rate should be 0.2
     model = _gen_efficientnet_lite(
@@ -1809,7 +1809,7 @@ def efficientnet_lite3(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_lite4(pretrained=False, **kwargs):
+def efficientnet_lite4(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-Lite4 """
     # NOTE for train, drop_rate should be 0.4, drop_path_rate should be 0.2
     model = _gen_efficientnet_lite(
@@ -1818,7 +1818,7 @@ def efficientnet_lite4(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_b1_pruned(pretrained=False, **kwargs):
+def efficientnet_b1_pruned(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B1 Pruned. The pruning has been obtained using https://arxiv.org/pdf/2002.08258.pdf  """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -1829,7 +1829,7 @@ def efficientnet_b1_pruned(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_b2_pruned(pretrained=False, **kwargs):
+def efficientnet_b2_pruned(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B2 Pruned. The pruning has been obtained using https://arxiv.org/pdf/2002.08258.pdf """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -1840,7 +1840,7 @@ def efficientnet_b2_pruned(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnet_b3_pruned(pretrained=False, **kwargs):
+def efficientnet_b3_pruned(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B3 Pruned. The pruning has been obtained using https://arxiv.org/pdf/2002.08258.pdf """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -1851,7 +1851,7 @@ def efficientnet_b3_pruned(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnetv2_rw_t(pretrained=False, **kwargs):
+def efficientnetv2_rw_t(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-V2 Tiny (Custom variant, tiny not in paper). """
     model = _gen_efficientnetv2_s(
         'efficientnetv2_rw_t', channel_multiplier=0.8, depth_multiplier=0.9, rw=False, pretrained=pretrained, **kwargs)
@@ -1859,7 +1859,7 @@ def efficientnetv2_rw_t(pretrained=False, **kwargs):
 
 
 @register_model
-def gc_efficientnetv2_rw_t(pretrained=False, **kwargs):
+def gc_efficientnetv2_rw_t(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-V2 Tiny w/ Global Context Attn (Custom variant, tiny not in paper). """
     model = _gen_efficientnetv2_s(
         'gc_efficientnetv2_rw_t', channel_multiplier=0.8, depth_multiplier=0.9,
@@ -1868,7 +1868,7 @@ def gc_efficientnetv2_rw_t(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnetv2_rw_s(pretrained=False, **kwargs):
+def efficientnetv2_rw_s(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-V2 Small (RW variant).
     NOTE: This is my initial (pre official code release) w/ some differences.
     See efficientnetv2_s and tf_efficientnetv2_s for versions that match the official w/ PyTorch vs TF padding
@@ -1878,7 +1878,7 @@ def efficientnetv2_rw_s(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnetv2_rw_m(pretrained=False, **kwargs):
+def efficientnetv2_rw_m(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-V2 Medium (RW variant).
     """
     model = _gen_efficientnetv2_s(
@@ -1888,35 +1888,35 @@ def efficientnetv2_rw_m(pretrained=False, **kwargs):
 
 
 @register_model
-def efficientnetv2_s(pretrained=False, **kwargs):
+def efficientnetv2_s(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-V2 Small. """
     model = _gen_efficientnetv2_s('efficientnetv2_s', pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def efficientnetv2_m(pretrained=False, **kwargs):
+def efficientnetv2_m(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-V2 Medium. """
     model = _gen_efficientnetv2_m('efficientnetv2_m', pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def efficientnetv2_l(pretrained=False, **kwargs):
+def efficientnetv2_l(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-V2 Large. """
     model = _gen_efficientnetv2_l('efficientnetv2_l', pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def efficientnetv2_xl(pretrained=False, **kwargs):
+def efficientnetv2_xl(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-V2 Xtra-Large. """
     model = _gen_efficientnetv2_xl('efficientnetv2_xl', pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def tf_efficientnet_b0(pretrained=False, **kwargs):
+def tf_efficientnet_b0(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B0. Tensorflow compatible variant  """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -1926,7 +1926,7 @@ def tf_efficientnet_b0(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnet_b1(pretrained=False, **kwargs):
+def tf_efficientnet_b1(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B1. Tensorflow compatible variant  """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -1936,7 +1936,7 @@ def tf_efficientnet_b1(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnet_b2(pretrained=False, **kwargs):
+def tf_efficientnet_b2(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B2. Tensorflow compatible variant  """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -1946,7 +1946,7 @@ def tf_efficientnet_b2(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnet_b3(pretrained=False, **kwargs):
+def tf_efficientnet_b3(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B3. Tensorflow compatible variant """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -1956,7 +1956,7 @@ def tf_efficientnet_b3(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnet_b4(pretrained=False, **kwargs):
+def tf_efficientnet_b4(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B4. Tensorflow compatible variant """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -1966,7 +1966,7 @@ def tf_efficientnet_b4(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnet_b5(pretrained=False, **kwargs):
+def tf_efficientnet_b5(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B5. Tensorflow compatible variant """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -1976,7 +1976,7 @@ def tf_efficientnet_b5(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnet_b6(pretrained=False, **kwargs):
+def tf_efficientnet_b6(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B6. Tensorflow compatible variant """
     # NOTE for train, drop_rate should be 0.5
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
@@ -1987,7 +1987,7 @@ def tf_efficientnet_b6(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnet_b7(pretrained=False, **kwargs):
+def tf_efficientnet_b7(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B7. Tensorflow compatible variant """
     # NOTE for train, drop_rate should be 0.5
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
@@ -1998,7 +1998,7 @@ def tf_efficientnet_b7(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnet_b8(pretrained=False, **kwargs):
+def tf_efficientnet_b8(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-B8. Tensorflow compatible variant """
     # NOTE for train, drop_rate should be 0.5
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
@@ -2009,7 +2009,7 @@ def tf_efficientnet_b8(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnet_l2(pretrained=False, **kwargs):
+def tf_efficientnet_l2(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-L2 NoisyStudent. Tensorflow compatible variant """
     # NOTE for train, drop_rate should be 0.5
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
@@ -2020,7 +2020,7 @@ def tf_efficientnet_l2(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnet_es(pretrained=False, **kwargs):
+def tf_efficientnet_es(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-Edge Small. Tensorflow compatible variant  """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -2030,7 +2030,7 @@ def tf_efficientnet_es(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnet_em(pretrained=False, **kwargs):
+def tf_efficientnet_em(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-Edge-Medium. Tensorflow compatible variant  """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -2040,7 +2040,7 @@ def tf_efficientnet_em(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnet_el(pretrained=False, **kwargs):
+def tf_efficientnet_el(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-Edge-Large. Tensorflow compatible variant  """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -2050,7 +2050,7 @@ def tf_efficientnet_el(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnet_cc_b0_4e(pretrained=False, **kwargs):
+def tf_efficientnet_cc_b0_4e(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-CondConv-B0 w/ 4 Experts. Tensorflow compatible variant """
     # NOTE for train, drop_rate should be 0.2, drop_path_rate should be 0.2
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
@@ -2061,7 +2061,7 @@ def tf_efficientnet_cc_b0_4e(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnet_cc_b0_8e(pretrained=False, **kwargs):
+def tf_efficientnet_cc_b0_8e(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-CondConv-B0 w/ 8 Experts. Tensorflow compatible variant """
     # NOTE for train, drop_rate should be 0.2, drop_path_rate should be 0.2
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
@@ -2073,7 +2073,7 @@ def tf_efficientnet_cc_b0_8e(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnet_cc_b1_8e(pretrained=False, **kwargs):
+def tf_efficientnet_cc_b1_8e(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-CondConv-B1 w/ 8 Experts. Tensorflow compatible variant """
     # NOTE for train, drop_rate should be 0.2, drop_path_rate should be 0.2
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
@@ -2085,7 +2085,7 @@ def tf_efficientnet_cc_b1_8e(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnet_lite0(pretrained=False, **kwargs):
+def tf_efficientnet_lite0(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-Lite0 """
     # NOTE for train, drop_rate should be 0.2, drop_path_rate should be 0.2
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
@@ -2096,7 +2096,7 @@ def tf_efficientnet_lite0(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnet_lite1(pretrained=False, **kwargs):
+def tf_efficientnet_lite1(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-Lite1 """
     # NOTE for train, drop_rate should be 0.2, drop_path_rate should be 0.2
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
@@ -2107,7 +2107,7 @@ def tf_efficientnet_lite1(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnet_lite2(pretrained=False, **kwargs):
+def tf_efficientnet_lite2(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-Lite2 """
     # NOTE for train, drop_rate should be 0.3, drop_path_rate should be 0.2
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
@@ -2118,7 +2118,7 @@ def tf_efficientnet_lite2(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnet_lite3(pretrained=False, **kwargs):
+def tf_efficientnet_lite3(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-Lite3 """
     # NOTE for train, drop_rate should be 0.3, drop_path_rate should be 0.2
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
@@ -2129,7 +2129,7 @@ def tf_efficientnet_lite3(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnet_lite4(pretrained=False, **kwargs):
+def tf_efficientnet_lite4(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-Lite4 """
     # NOTE for train, drop_rate should be 0.4, drop_path_rate should be 0.2
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
@@ -2140,7 +2140,7 @@ def tf_efficientnet_lite4(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnetv2_s(pretrained=False, **kwargs):
+def tf_efficientnetv2_s(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-V2 Small. Tensorflow compatible variant  """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -2149,7 +2149,7 @@ def tf_efficientnetv2_s(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnetv2_m(pretrained=False, **kwargs):
+def tf_efficientnetv2_m(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-V2 Medium. Tensorflow compatible variant  """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -2158,7 +2158,7 @@ def tf_efficientnetv2_m(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnetv2_l(pretrained=False, **kwargs):
+def tf_efficientnetv2_l(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-V2 Large. Tensorflow compatible variant  """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -2167,7 +2167,7 @@ def tf_efficientnetv2_l(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnetv2_xl(pretrained=False, **kwargs):
+def tf_efficientnetv2_xl(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-V2 Xtra-Large. Tensorflow compatible variant
     """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
@@ -2177,7 +2177,7 @@ def tf_efficientnetv2_xl(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnetv2_b0(pretrained=False, **kwargs):
+def tf_efficientnetv2_b0(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-V2-B0. Tensorflow compatible variant  """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -2186,7 +2186,7 @@ def tf_efficientnetv2_b0(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnetv2_b1(pretrained=False, **kwargs):
+def tf_efficientnetv2_b1(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-V2-B1. Tensorflow compatible variant  """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -2196,7 +2196,7 @@ def tf_efficientnetv2_b1(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnetv2_b2(pretrained=False, **kwargs):
+def tf_efficientnetv2_b2(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-V2-B2. Tensorflow compatible variant  """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -2206,7 +2206,7 @@ def tf_efficientnetv2_b2(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_efficientnetv2_b3(pretrained=False, **kwargs):
+def tf_efficientnetv2_b3(pretrained=False, **kwargs) -> EfficientNet:
     """ EfficientNet-V2-B3. Tensorflow compatible variant """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -2216,7 +2216,7 @@ def tf_efficientnetv2_b3(pretrained=False, **kwargs):
 
 
 @register_model
-def mixnet_s(pretrained=False, **kwargs):
+def mixnet_s(pretrained=False, **kwargs) -> EfficientNet:
     """Creates a MixNet Small model.
     """
     model = _gen_mixnet_s(
@@ -2225,7 +2225,7 @@ def mixnet_s(pretrained=False, **kwargs):
 
 
 @register_model
-def mixnet_m(pretrained=False, **kwargs):
+def mixnet_m(pretrained=False, **kwargs) -> EfficientNet:
     """Creates a MixNet Medium model.
     """
     model = _gen_mixnet_m(
@@ -2234,7 +2234,7 @@ def mixnet_m(pretrained=False, **kwargs):
 
 
 @register_model
-def mixnet_l(pretrained=False, **kwargs):
+def mixnet_l(pretrained=False, **kwargs) -> EfficientNet:
     """Creates a MixNet Large model.
     """
     model = _gen_mixnet_m(
@@ -2243,7 +2243,7 @@ def mixnet_l(pretrained=False, **kwargs):
 
 
 @register_model
-def mixnet_xl(pretrained=False, **kwargs):
+def mixnet_xl(pretrained=False, **kwargs) -> EfficientNet:
     """Creates a MixNet Extra-Large model.
     Not a paper spec, experimental def by RW w/ depth scaling.
     """
@@ -2253,7 +2253,7 @@ def mixnet_xl(pretrained=False, **kwargs):
 
 
 @register_model
-def mixnet_xxl(pretrained=False, **kwargs):
+def mixnet_xxl(pretrained=False, **kwargs) -> EfficientNet:
     """Creates a MixNet Double Extra Large model.
     Not a paper spec, experimental def by RW w/ depth scaling.
     """
@@ -2263,7 +2263,7 @@ def mixnet_xxl(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_mixnet_s(pretrained=False, **kwargs):
+def tf_mixnet_s(pretrained=False, **kwargs) -> EfficientNet:
     """Creates a MixNet Small model. Tensorflow compatible variant
     """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
@@ -2274,7 +2274,7 @@ def tf_mixnet_s(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_mixnet_m(pretrained=False, **kwargs):
+def tf_mixnet_m(pretrained=False, **kwargs) -> EfficientNet:
     """Creates a MixNet Medium model. Tensorflow compatible variant
     """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
@@ -2285,7 +2285,7 @@ def tf_mixnet_m(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_mixnet_l(pretrained=False, **kwargs):
+def tf_mixnet_l(pretrained=False, **kwargs) -> EfficientNet:
     """Creates a MixNet Large model. Tensorflow compatible variant
     """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
@@ -2296,31 +2296,31 @@ def tf_mixnet_l(pretrained=False, **kwargs):
 
 
 @register_model
-def tinynet_a(pretrained=False, **kwargs):
+def tinynet_a(pretrained=False, **kwargs) -> EfficientNet:
     model = _gen_tinynet('tinynet_a', 1.0, 1.2, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def tinynet_b(pretrained=False, **kwargs):
+def tinynet_b(pretrained=False, **kwargs) -> EfficientNet:
     model = _gen_tinynet('tinynet_b', 0.75, 1.1, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def tinynet_c(pretrained=False, **kwargs):
+def tinynet_c(pretrained=False, **kwargs) -> EfficientNet:
     model = _gen_tinynet('tinynet_c', 0.54, 0.85, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def tinynet_d(pretrained=False, **kwargs):
+def tinynet_d(pretrained=False, **kwargs) -> EfficientNet:
     model = _gen_tinynet('tinynet_d', 0.54, 0.695, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def tinynet_e(pretrained=False, **kwargs):
+def tinynet_e(pretrained=False, **kwargs) -> EfficientNet:
     model = _gen_tinynet('tinynet_e', 0.51, 0.6, pretrained=pretrained, **kwargs)
     return model
 
diff --git a/timm/models/eva.py b/timm/models/eva.py
index 6142b2ded2..6c01d447d5 100644
--- a/timm/models/eva.py
+++ b/timm/models/eva.py
@@ -825,7 +825,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def eva_giant_patch14_224(pretrained=False, **kwargs):
+def eva_giant_patch14_224(pretrained=False, **kwargs) -> Eva:
     """ EVA-g model https://arxiv.org/abs/2211.07636 """
     model_args = dict(patch_size=14, embed_dim=1408, depth=40, num_heads=16, mlp_ratio=6144 / 1408)
     model = _create_eva('eva_giant_patch14_224', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -833,7 +833,7 @@ def eva_giant_patch14_224(pretrained=False, **kwargs):
 
 
 @register_model
-def eva_giant_patch14_336(pretrained=False, **kwargs):
+def eva_giant_patch14_336(pretrained=False, **kwargs) -> Eva:
     """ EVA-g model https://arxiv.org/abs/2211.07636 """
     model_args = dict(patch_size=14, embed_dim=1408, depth=40, num_heads=16, mlp_ratio=6144 / 1408)
     model = _create_eva('eva_giant_patch14_336', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -841,7 +841,7 @@ def eva_giant_patch14_336(pretrained=False, **kwargs):
 
 
 @register_model
-def eva_giant_patch14_560(pretrained=False, **kwargs):
+def eva_giant_patch14_560(pretrained=False, **kwargs) -> Eva:
     """ EVA-g model https://arxiv.org/abs/2211.07636 """
     model_args = dict(patch_size=14, embed_dim=1408, depth=40, num_heads=16, mlp_ratio=6144 / 1408)
     model = _create_eva('eva_giant_patch14_560', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -849,7 +849,7 @@ def eva_giant_patch14_560(pretrained=False, **kwargs):
 
 
 @register_model
-def eva02_tiny_patch14_224(pretrained=False, **kwargs):
+def eva02_tiny_patch14_224(pretrained=False, **kwargs) -> Eva:
     model_args = dict(
         img_size=224,
         patch_size=14,
@@ -866,7 +866,7 @@ def eva02_tiny_patch14_224(pretrained=False, **kwargs):
 
 
 @register_model
-def eva02_small_patch14_224(pretrained=False, **kwargs):
+def eva02_small_patch14_224(pretrained=False, **kwargs) -> Eva:
     model_args = dict(
         img_size=224,
         patch_size=14,
@@ -883,7 +883,7 @@ def eva02_small_patch14_224(pretrained=False, **kwargs):
 
 
 @register_model
-def eva02_base_patch14_224(pretrained=False, **kwargs):
+def eva02_base_patch14_224(pretrained=False, **kwargs) -> Eva:
     model_args = dict(
         img_size=224,
         patch_size=14,
@@ -902,7 +902,7 @@ def eva02_base_patch14_224(pretrained=False, **kwargs):
 
 
 @register_model
-def eva02_large_patch14_224(pretrained=False, **kwargs):
+def eva02_large_patch14_224(pretrained=False, **kwargs) -> Eva:
     model_args = dict(
         img_size=224,
         patch_size=14,
@@ -921,7 +921,7 @@ def eva02_large_patch14_224(pretrained=False, **kwargs):
 
 
 @register_model
-def eva02_tiny_patch14_336(pretrained=False, **kwargs):
+def eva02_tiny_patch14_336(pretrained=False, **kwargs) -> Eva:
     model_args = dict(
         img_size=336,
         patch_size=14,
@@ -938,7 +938,7 @@ def eva02_tiny_patch14_336(pretrained=False, **kwargs):
 
 
 @register_model
-def eva02_small_patch14_336(pretrained=False, **kwargs):
+def eva02_small_patch14_336(pretrained=False, **kwargs) -> Eva:
     model_args = dict(
         img_size=336,
         patch_size=14,
@@ -955,7 +955,7 @@ def eva02_small_patch14_336(pretrained=False, **kwargs):
 
 
 @register_model
-def eva02_base_patch14_448(pretrained=False, **kwargs):
+def eva02_base_patch14_448(pretrained=False, **kwargs) -> Eva:
     model_args = dict(
         img_size=448,
         patch_size=14,
@@ -974,7 +974,7 @@ def eva02_base_patch14_448(pretrained=False, **kwargs):
 
 
 @register_model
-def eva02_large_patch14_448(pretrained=False, **kwargs):
+def eva02_large_patch14_448(pretrained=False, **kwargs) -> Eva:
     model_args = dict(
         img_size=448,
         patch_size=14,
@@ -993,7 +993,7 @@ def eva02_large_patch14_448(pretrained=False, **kwargs):
 
 
 @register_model
-def eva_giant_patch14_clip_224(pretrained=False, **kwargs):
+def eva_giant_patch14_clip_224(pretrained=False, **kwargs) -> Eva:
     """ EVA-g CLIP model (only difference from non-CLIP is the pooling)  """
     model_args = dict(
         patch_size=14, embed_dim=1408, depth=40, num_heads=16, mlp_ratio=6144 / 1408,
@@ -1003,7 +1003,7 @@ def eva_giant_patch14_clip_224(pretrained=False, **kwargs):
 
 
 @register_model
-def eva02_base_patch16_clip_224(pretrained=False, **kwargs):
+def eva02_base_patch16_clip_224(pretrained=False, **kwargs) -> Eva:
     """ A EVA-CLIP specific variant that adds additional attn scale layernorm to eva02_base """
     model_args = dict(
         img_size=224,
@@ -1025,7 +1025,7 @@ def eva02_base_patch16_clip_224(pretrained=False, **kwargs):
 
 
 @register_model
-def eva02_large_patch14_clip_224(pretrained=False, **kwargs):
+def eva02_large_patch14_clip_224(pretrained=False, **kwargs) -> Eva:
     """ A EVA-CLIP specific variant that adds additional attn scale layernorm to eva02_large """
     model_args = dict(
         img_size=224,
@@ -1047,7 +1047,7 @@ def eva02_large_patch14_clip_224(pretrained=False, **kwargs):
 
 
 @register_model
-def eva02_large_patch14_clip_336(pretrained=False, **kwargs):
+def eva02_large_patch14_clip_336(pretrained=False, **kwargs) -> Eva:
     """ A EVA-CLIP specific variant that adds additional attn scale layernorm to eva02_large """
     model_args = dict(
         img_size=336,
@@ -1069,7 +1069,7 @@ def eva02_large_patch14_clip_336(pretrained=False, **kwargs):
 
 
 @register_model
-def eva02_enormous_patch14_clip_224(pretrained=False, **kwargs):
+def eva02_enormous_patch14_clip_224(pretrained=False, **kwargs) -> Eva:
     """ A EVA-CLIP specific variant that uses residual post-norm in blocks """
     model_args = dict(
         img_size=224,
diff --git a/timm/models/focalnet.py b/timm/models/focalnet.py
index 104b09fcec..b8d90db0eb 100644
--- a/timm/models/focalnet.py
+++ b/timm/models/focalnet.py
@@ -566,44 +566,44 @@ def _create_focalnet(variant, pretrained=False, **kwargs):
 
 
 @register_model
-def focalnet_tiny_srf(pretrained=False, **kwargs):
+def focalnet_tiny_srf(pretrained=False, **kwargs) -> FocalNet:
     model_kwargs = dict(depths=[2, 2, 6, 2], embed_dim=96, **kwargs)
     return _create_focalnet('focalnet_tiny_srf', pretrained=pretrained, **model_kwargs)
 
 
 @register_model
-def focalnet_small_srf(pretrained=False, **kwargs):
+def focalnet_small_srf(pretrained=False, **kwargs) -> FocalNet:
     model_kwargs = dict(depths=[2, 2, 18, 2], embed_dim=96, **kwargs)
     return _create_focalnet('focalnet_small_srf', pretrained=pretrained, **model_kwargs)
 
 
 @register_model
-def focalnet_base_srf(pretrained=False, **kwargs):
+def focalnet_base_srf(pretrained=False, **kwargs) -> FocalNet:
     model_kwargs = dict(depths=[2, 2, 18, 2], embed_dim=128, **kwargs)
     return _create_focalnet('focalnet_base_srf', pretrained=pretrained, **model_kwargs)
 
 
 @register_model
-def focalnet_tiny_lrf(pretrained=False, **kwargs):
+def focalnet_tiny_lrf(pretrained=False, **kwargs) -> FocalNet:
     model_kwargs = dict(depths=[2, 2, 6, 2], embed_dim=96, focal_levels=[3, 3, 3, 3], **kwargs)
     return _create_focalnet('focalnet_tiny_lrf', pretrained=pretrained, **model_kwargs)
 
 
 @register_model
-def focalnet_small_lrf(pretrained=False, **kwargs):
+def focalnet_small_lrf(pretrained=False, **kwargs) -> FocalNet:
     model_kwargs = dict(depths=[2, 2, 18, 2], embed_dim=96, focal_levels=[3, 3, 3, 3], **kwargs)
     return _create_focalnet('focalnet_small_lrf', pretrained=pretrained, **model_kwargs)
 
 
 @register_model
-def focalnet_base_lrf(pretrained=False, **kwargs):
+def focalnet_base_lrf(pretrained=False, **kwargs) -> FocalNet:
     model_kwargs = dict(depths=[2, 2, 18, 2], embed_dim=128, focal_levels=[3, 3, 3, 3], **kwargs)
     return _create_focalnet('focalnet_base_lrf', pretrained=pretrained, **model_kwargs)
 
 
 # FocalNet large+ models
 @register_model
-def focalnet_large_fl3(pretrained=False, **kwargs):
+def focalnet_large_fl3(pretrained=False, **kwargs) -> FocalNet:
     model_kwargs = dict(
         depths=[2, 2, 18, 2], embed_dim=192, focal_levels=[3, 3, 3, 3], focal_windows=[5] * 4,
         use_post_norm=True, use_overlap_down=True, layerscale_value=1e-4, **kwargs)
@@ -611,7 +611,7 @@ def focalnet_large_fl3(pretrained=False, **kwargs):
 
 
 @register_model
-def focalnet_large_fl4(pretrained=False, **kwargs):
+def focalnet_large_fl4(pretrained=False, **kwargs) -> FocalNet:
     model_kwargs = dict(
         depths=[2, 2, 18, 2], embed_dim=192, focal_levels=[4, 4, 4, 4],
         use_post_norm=True, use_overlap_down=True, layerscale_value=1e-4, **kwargs)
@@ -619,7 +619,7 @@ def focalnet_large_fl4(pretrained=False, **kwargs):
 
 
 @register_model
-def focalnet_xlarge_fl3(pretrained=False, **kwargs):
+def focalnet_xlarge_fl3(pretrained=False, **kwargs) -> FocalNet:
     model_kwargs = dict(
         depths=[2, 2, 18, 2], embed_dim=256, focal_levels=[3, 3, 3, 3], focal_windows=[5] * 4,
         use_post_norm=True, use_overlap_down=True, layerscale_value=1e-4, **kwargs)
@@ -627,7 +627,7 @@ def focalnet_xlarge_fl3(pretrained=False, **kwargs):
 
 
 @register_model
-def focalnet_xlarge_fl4(pretrained=False, **kwargs):
+def focalnet_xlarge_fl4(pretrained=False, **kwargs) -> FocalNet:
     model_kwargs = dict(
         depths=[2, 2, 18, 2], embed_dim=256, focal_levels=[4, 4, 4, 4],
         use_post_norm=True, use_overlap_down=True, layerscale_value=1e-4, **kwargs)
@@ -635,7 +635,7 @@ def focalnet_xlarge_fl4(pretrained=False, **kwargs):
 
 
 @register_model
-def focalnet_huge_fl3(pretrained=False, **kwargs):
+def focalnet_huge_fl3(pretrained=False, **kwargs) -> FocalNet:
     model_kwargs = dict(
         depths=[2, 2, 18, 2], embed_dim=352, focal_levels=[3, 3, 3, 3], focal_windows=[3] * 4,
         use_post_norm=True, use_post_norm_in_modulation=True, use_overlap_down=True, layerscale_value=1e-4, **kwargs)
@@ -643,7 +643,7 @@ def focalnet_huge_fl3(pretrained=False, **kwargs):
 
 
 @register_model
-def focalnet_huge_fl4(pretrained=False, **kwargs):
+def focalnet_huge_fl4(pretrained=False, **kwargs) -> FocalNet:
     model_kwargs = dict(
         depths=[2, 2, 18, 2], embed_dim=352, focal_levels=[4, 4, 4, 4],
         use_post_norm=True, use_post_norm_in_modulation=True, use_overlap_down=True, layerscale_value=1e-4, **kwargs)
diff --git a/timm/models/gcvit.py b/timm/models/gcvit.py
index 32d2e7035d..29cdb1fbe0 100644
--- a/timm/models/gcvit.py
+++ b/timm/models/gcvit.py
@@ -542,7 +542,7 @@ def _create_gcvit(variant, pretrained=False, **kwargs):
 
 
 @register_model
-def gcvit_xxtiny(pretrained=False, **kwargs):
+def gcvit_xxtiny(pretrained=False, **kwargs) -> GlobalContextVit:
     model_kwargs = dict(
         depths=(2, 2, 6, 2),
         num_heads=(2, 4, 8, 16),
@@ -551,7 +551,7 @@ def gcvit_xxtiny(pretrained=False, **kwargs):
 
 
 @register_model
-def gcvit_xtiny(pretrained=False, **kwargs):
+def gcvit_xtiny(pretrained=False, **kwargs) -> GlobalContextVit:
     model_kwargs = dict(
         depths=(3, 4, 6, 5),
         num_heads=(2, 4, 8, 16),
@@ -560,7 +560,7 @@ def gcvit_xtiny(pretrained=False, **kwargs):
 
 
 @register_model
-def gcvit_tiny(pretrained=False, **kwargs):
+def gcvit_tiny(pretrained=False, **kwargs) -> GlobalContextVit:
     model_kwargs = dict(
         depths=(3, 4, 19, 5),
         num_heads=(2, 4, 8, 16),
@@ -569,7 +569,7 @@ def gcvit_tiny(pretrained=False, **kwargs):
 
 
 @register_model
-def gcvit_small(pretrained=False, **kwargs):
+def gcvit_small(pretrained=False, **kwargs) -> GlobalContextVit:
     model_kwargs = dict(
         depths=(3, 4, 19, 5),
         num_heads=(3, 6, 12, 24),
@@ -581,7 +581,7 @@ def gcvit_small(pretrained=False, **kwargs):
 
 
 @register_model
-def gcvit_base(pretrained=False, **kwargs):
+def gcvit_base(pretrained=False, **kwargs) -> GlobalContextVit:
     model_kwargs = dict(
         depths=(3, 4, 19, 5),
         num_heads=(4, 8, 16, 32),
diff --git a/timm/models/ghostnet.py b/timm/models/ghostnet.py
index 4838ee6261..f5b5123f54 100644
--- a/timm/models/ghostnet.py
+++ b/timm/models/ghostnet.py
@@ -309,21 +309,21 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def ghostnet_050(pretrained=False, **kwargs):
+def ghostnet_050(pretrained=False, **kwargs) -> GhostNet:
     """ GhostNet-0.5x """
     model = _create_ghostnet('ghostnet_050', width=0.5, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def ghostnet_100(pretrained=False, **kwargs):
+def ghostnet_100(pretrained=False, **kwargs) -> GhostNet:
     """ GhostNet-1.0x """
     model = _create_ghostnet('ghostnet_100', width=1.0, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def ghostnet_130(pretrained=False, **kwargs):
+def ghostnet_130(pretrained=False, **kwargs) -> GhostNet:
     """ GhostNet-1.3x """
     model = _create_ghostnet('ghostnet_130', width=1.3, pretrained=pretrained, **kwargs)
     return model
diff --git a/timm/models/hardcorenas.py b/timm/models/hardcorenas.py
index d3835a2460..459c1a3db8 100644
--- a/timm/models/hardcorenas.py
+++ b/timm/models/hardcorenas.py
@@ -73,7 +73,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def hardcorenas_a(pretrained=False, **kwargs):
+def hardcorenas_a(pretrained=False, **kwargs) -> MobileNetV3:
     """ hardcorenas_A """
     arch_def = [['ds_r1_k3_s1_e1_c16_nre'], ['ir_r1_k5_s2_e3_c24_nre', 'ir_r1_k5_s1_e3_c24_nre_se0.25'],
                 ['ir_r1_k5_s2_e3_c40_nre', 'ir_r1_k5_s1_e6_c40_nre_se0.25'],
@@ -85,7 +85,7 @@ def hardcorenas_a(pretrained=False, **kwargs):
 
 
 @register_model
-def hardcorenas_b(pretrained=False, **kwargs):
+def hardcorenas_b(pretrained=False, **kwargs) -> MobileNetV3:
     """ hardcorenas_B """
     arch_def = [['ds_r1_k3_s1_e1_c16_nre'],
                 ['ir_r1_k5_s2_e3_c24_nre', 'ir_r1_k5_s1_e3_c24_nre_se0.25', 'ir_r1_k3_s1_e3_c24_nre'],
@@ -99,7 +99,7 @@ def hardcorenas_b(pretrained=False, **kwargs):
 
 
 @register_model
-def hardcorenas_c(pretrained=False, **kwargs):
+def hardcorenas_c(pretrained=False, **kwargs) -> MobileNetV3:
     """ hardcorenas_C """
     arch_def = [['ds_r1_k3_s1_e1_c16_nre'], ['ir_r1_k5_s2_e3_c24_nre', 'ir_r1_k5_s1_e3_c24_nre_se0.25'],
                 ['ir_r1_k5_s2_e3_c40_nre', 'ir_r1_k5_s1_e3_c40_nre', 'ir_r1_k5_s1_e3_c40_nre',
@@ -113,7 +113,7 @@ def hardcorenas_c(pretrained=False, **kwargs):
 
 
 @register_model
-def hardcorenas_d(pretrained=False, **kwargs):
+def hardcorenas_d(pretrained=False, **kwargs) -> MobileNetV3:
     """ hardcorenas_D """
     arch_def = [['ds_r1_k3_s1_e1_c16_nre'], ['ir_r1_k5_s2_e3_c24_nre_se0.25', 'ir_r1_k5_s1_e3_c24_nre_se0.25'],
                 ['ir_r1_k5_s2_e3_c40_nre_se0.25', 'ir_r1_k5_s1_e4_c40_nre_se0.25', 'ir_r1_k3_s1_e3_c40_nre_se0.25'],
@@ -128,7 +128,7 @@ def hardcorenas_d(pretrained=False, **kwargs):
 
 
 @register_model
-def hardcorenas_e(pretrained=False, **kwargs):
+def hardcorenas_e(pretrained=False, **kwargs) -> MobileNetV3:
     """ hardcorenas_E """
     arch_def = [['ds_r1_k3_s1_e1_c16_nre'], ['ir_r1_k5_s2_e3_c24_nre_se0.25', 'ir_r1_k5_s1_e3_c24_nre_se0.25'],
                 ['ir_r1_k5_s2_e6_c40_nre_se0.25', 'ir_r1_k5_s1_e4_c40_nre_se0.25', 'ir_r1_k5_s1_e4_c40_nre_se0.25',
@@ -142,7 +142,7 @@ def hardcorenas_e(pretrained=False, **kwargs):
 
 
 @register_model
-def hardcorenas_f(pretrained=False, **kwargs):
+def hardcorenas_f(pretrained=False, **kwargs) -> MobileNetV3:
     """ hardcorenas_F """
     arch_def = [['ds_r1_k3_s1_e1_c16_nre'], ['ir_r1_k5_s2_e3_c24_nre_se0.25', 'ir_r1_k5_s1_e3_c24_nre_se0.25'],
                 ['ir_r1_k5_s2_e6_c40_nre_se0.25', 'ir_r1_k5_s1_e6_c40_nre_se0.25'],
diff --git a/timm/models/hrnet.py b/timm/models/hrnet.py
index db75bc0f3b..4a6fd80eaa 100644
--- a/timm/models/hrnet.py
+++ b/timm/models/hrnet.py
@@ -915,58 +915,58 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def hrnet_w18_small(pretrained=False, **kwargs):
+def hrnet_w18_small(pretrained=False, **kwargs) -> HighResolutionNet:
     return _create_hrnet('hrnet_w18_small', pretrained, **kwargs)
 
 
 @register_model
-def hrnet_w18_small_v2(pretrained=False, **kwargs):
+def hrnet_w18_small_v2(pretrained=False, **kwargs) -> HighResolutionNet:
     return _create_hrnet('hrnet_w18_small_v2', pretrained, **kwargs)
 
 
 @register_model
-def hrnet_w18(pretrained=False, **kwargs):
+def hrnet_w18(pretrained=False, **kwargs) -> HighResolutionNet:
     return _create_hrnet('hrnet_w18', pretrained, **kwargs)
 
 
 @register_model
-def hrnet_w30(pretrained=False, **kwargs):
+def hrnet_w30(pretrained=False, **kwargs) -> HighResolutionNet:
     return _create_hrnet('hrnet_w30', pretrained, **kwargs)
 
 
 @register_model
-def hrnet_w32(pretrained=False, **kwargs):
+def hrnet_w32(pretrained=False, **kwargs) -> HighResolutionNet:
     return _create_hrnet('hrnet_w32', pretrained, **kwargs)
 
 
 @register_model
-def hrnet_w40(pretrained=False, **kwargs):
+def hrnet_w40(pretrained=False, **kwargs) -> HighResolutionNet:
     return _create_hrnet('hrnet_w40', pretrained, **kwargs)
 
 
 @register_model
-def hrnet_w44(pretrained=False, **kwargs):
+def hrnet_w44(pretrained=False, **kwargs) -> HighResolutionNet:
     return _create_hrnet('hrnet_w44', pretrained, **kwargs)
 
 
 @register_model
-def hrnet_w48(pretrained=False, **kwargs):
+def hrnet_w48(pretrained=False, **kwargs) -> HighResolutionNet:
     return _create_hrnet('hrnet_w48', pretrained, **kwargs)
 
 
 @register_model
-def hrnet_w64(pretrained=False, **kwargs):
+def hrnet_w64(pretrained=False, **kwargs) -> HighResolutionNet:
     return _create_hrnet('hrnet_w64', pretrained, **kwargs)
 
 
 @register_model
-def hrnet_w18_ssld(pretrained=False, **kwargs):
+def hrnet_w18_ssld(pretrained=False, **kwargs) -> HighResolutionNet:
     kwargs.setdefault('head_conv_bias', False)
     return _create_hrnet('hrnet_w18_ssld', cfg_variant='hrnet_w18', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def hrnet_w48_ssld(pretrained=False, **kwargs):
+def hrnet_w48_ssld(pretrained=False, **kwargs) -> HighResolutionNet:
     kwargs.setdefault('head_conv_bias', False)
     return _create_hrnet('hrnet_w48_ssld', cfg_variant='hrnet_w48', pretrained=pretrained, **kwargs)
 
diff --git a/timm/models/inception_resnet_v2.py b/timm/models/inception_resnet_v2.py
index f69a1a8131..138470f750 100644
--- a/timm/models/inception_resnet_v2.py
+++ b/timm/models/inception_resnet_v2.py
@@ -336,7 +336,7 @@ def _create_inception_resnet_v2(variant, pretrained=False, **kwargs):
 
 
 @register_model
-def inception_resnet_v2(pretrained=False, **kwargs):
+def inception_resnet_v2(pretrained=False, **kwargs) -> InceptionResnetV2:
     return _create_inception_resnet_v2('inception_resnet_v2', pretrained=pretrained, **kwargs)
 
 
diff --git a/timm/models/inception_v3.py b/timm/models/inception_v3.py
index dc9e130248..984970817d 100644
--- a/timm/models/inception_v3.py
+++ b/timm/models/inception_v3.py
@@ -450,7 +450,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def inception_v3(pretrained=False, **kwargs):
+def inception_v3(pretrained=False, **kwargs) -> InceptionV3:
     model = _create_inception_v3('inception_v3', pretrained=pretrained, **kwargs)
     return model
 
diff --git a/timm/models/inception_v4.py b/timm/models/inception_v4.py
index e225a48f12..91c435d978 100644
--- a/timm/models/inception_v4.py
+++ b/timm/models/inception_v4.py
@@ -299,7 +299,7 @@ def forward(self, x):
         return x
 
 
-def _create_inception_v4(variant, pretrained=False, **kwargs):
+def _create_inception_v4(variant, pretrained=False, **kwargs) -> InceptionV4:
     return build_model_with_cfg(
         InceptionV4,
         variant,
diff --git a/timm/models/levit.py b/timm/models/levit.py
index 3ee3534c8a..71ce4a364f 100644
--- a/timm/models/levit.py
+++ b/timm/models/levit.py
@@ -832,101 +832,101 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def levit_128s(pretrained=False, **kwargs):
+def levit_128s(pretrained=False, **kwargs) -> Levit:
     return create_levit('levit_128s', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def levit_128(pretrained=False, **kwargs):
+def levit_128(pretrained=False, **kwargs) -> Levit:
     return create_levit('levit_128', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def levit_192(pretrained=False, **kwargs):
+def levit_192(pretrained=False, **kwargs) -> Levit:
     return create_levit('levit_192', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def levit_256(pretrained=False, **kwargs):
+def levit_256(pretrained=False, **kwargs) -> Levit:
     return create_levit('levit_256', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def levit_384(pretrained=False, **kwargs):
+def levit_384(pretrained=False, **kwargs) -> Levit:
     return create_levit('levit_384', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def levit_384_s8(pretrained=False, **kwargs):
+def levit_384_s8(pretrained=False, **kwargs) -> Levit:
     return create_levit('levit_384_s8', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def levit_512_s8(pretrained=False, **kwargs):
+def levit_512_s8(pretrained=False, **kwargs) -> Levit:
     return create_levit('levit_512_s8', pretrained=pretrained, distilled=False, **kwargs)
 
 
 @register_model
-def levit_512(pretrained=False, **kwargs):
+def levit_512(pretrained=False, **kwargs) -> Levit:
     return create_levit('levit_512', pretrained=pretrained, distilled=False, **kwargs)
 
 
 @register_model
-def levit_256d(pretrained=False, **kwargs):
+def levit_256d(pretrained=False, **kwargs) -> Levit:
     return create_levit('levit_256d', pretrained=pretrained, distilled=False, **kwargs)
 
 
 @register_model
-def levit_512d(pretrained=False, **kwargs):
+def levit_512d(pretrained=False, **kwargs) -> Levit:
     return create_levit('levit_512d', pretrained=pretrained, distilled=False, **kwargs)
 
 
 @register_model
-def levit_conv_128s(pretrained=False, **kwargs):
+def levit_conv_128s(pretrained=False, **kwargs) -> Levit:
     return create_levit('levit_conv_128s', pretrained=pretrained, use_conv=True, **kwargs)
 
 
 @register_model
-def levit_conv_128(pretrained=False, **kwargs):
+def levit_conv_128(pretrained=False, **kwargs) -> Levit:
     return create_levit('levit_conv_128', pretrained=pretrained, use_conv=True, **kwargs)
 
 
 @register_model
-def levit_conv_192(pretrained=False, **kwargs):
+def levit_conv_192(pretrained=False, **kwargs) -> Levit:
     return create_levit('levit_conv_192', pretrained=pretrained, use_conv=True, **kwargs)
 
 
 @register_model
-def levit_conv_256(pretrained=False, **kwargs):
+def levit_conv_256(pretrained=False, **kwargs) -> Levit:
     return create_levit('levit_conv_256', pretrained=pretrained, use_conv=True, **kwargs)
 
 
 @register_model
-def levit_conv_384(pretrained=False, **kwargs):
+def levit_conv_384(pretrained=False, **kwargs) -> Levit:
     return create_levit('levit_conv_384', pretrained=pretrained, use_conv=True, **kwargs)
 
 
 @register_model
-def levit_conv_384_s8(pretrained=False, **kwargs):
+def levit_conv_384_s8(pretrained=False, **kwargs) -> Levit:
     return create_levit('levit_conv_384_s8', pretrained=pretrained, use_conv=True, **kwargs)
 
 
 @register_model
-def levit_conv_512_s8(pretrained=False, **kwargs):
+def levit_conv_512_s8(pretrained=False, **kwargs) -> Levit:
     return create_levit('levit_conv_512_s8', pretrained=pretrained, use_conv=True, distilled=False, **kwargs)
 
 
 @register_model
-def levit_conv_512(pretrained=False, **kwargs):
+def levit_conv_512(pretrained=False, **kwargs) -> Levit:
     return create_levit('levit_conv_512', pretrained=pretrained, use_conv=True, distilled=False, **kwargs)
 
 
 @register_model
-def levit_conv_256d(pretrained=False, **kwargs):
+def levit_conv_256d(pretrained=False, **kwargs) -> Levit:
     return create_levit('levit_conv_256d', pretrained=pretrained, use_conv=True, distilled=False, **kwargs)
 
 
 @register_model
-def levit_conv_512d(pretrained=False, **kwargs):
+def levit_conv_512d(pretrained=False, **kwargs) -> Levit:
     return create_levit('levit_conv_512d', pretrained=pretrained, use_conv=True, distilled=False, **kwargs)
 
diff --git a/timm/models/maxxvit.py b/timm/models/maxxvit.py
index 17b64a9dbb..76ca18a0e5 100644
--- a/timm/models/maxxvit.py
+++ b/timm/models/maxxvit.py
@@ -2040,280 +2040,280 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def coatnet_pico_rw_224(pretrained=False, **kwargs):
+def coatnet_pico_rw_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('coatnet_pico_rw_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def coatnet_nano_rw_224(pretrained=False, **kwargs):
+def coatnet_nano_rw_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('coatnet_nano_rw_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def coatnet_0_rw_224(pretrained=False, **kwargs):
+def coatnet_0_rw_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('coatnet_0_rw_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def coatnet_1_rw_224(pretrained=False, **kwargs):
+def coatnet_1_rw_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('coatnet_1_rw_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def coatnet_2_rw_224(pretrained=False, **kwargs):
+def coatnet_2_rw_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('coatnet_2_rw_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def coatnet_3_rw_224(pretrained=False, **kwargs):
+def coatnet_3_rw_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('coatnet_3_rw_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def coatnet_bn_0_rw_224(pretrained=False, **kwargs):
+def coatnet_bn_0_rw_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('coatnet_bn_0_rw_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def coatnet_rmlp_nano_rw_224(pretrained=False, **kwargs):
+def coatnet_rmlp_nano_rw_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('coatnet_rmlp_nano_rw_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def coatnet_rmlp_0_rw_224(pretrained=False, **kwargs):
+def coatnet_rmlp_0_rw_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('coatnet_rmlp_0_rw_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def coatnet_rmlp_1_rw_224(pretrained=False, **kwargs):
+def coatnet_rmlp_1_rw_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('coatnet_rmlp_1_rw_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def coatnet_rmlp_1_rw2_224(pretrained=False, **kwargs):
+def coatnet_rmlp_1_rw2_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('coatnet_rmlp_1_rw2_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def coatnet_rmlp_2_rw_224(pretrained=False, **kwargs):
+def coatnet_rmlp_2_rw_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('coatnet_rmlp_2_rw_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def coatnet_rmlp_2_rw_384(pretrained=False, **kwargs):
+def coatnet_rmlp_2_rw_384(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('coatnet_rmlp_2_rw_384', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def coatnet_rmlp_3_rw_224(pretrained=False, **kwargs):
+def coatnet_rmlp_3_rw_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('coatnet_rmlp_3_rw_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def coatnet_nano_cc_224(pretrained=False, **kwargs):
+def coatnet_nano_cc_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('coatnet_nano_cc_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def coatnext_nano_rw_224(pretrained=False, **kwargs):
+def coatnext_nano_rw_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('coatnext_nano_rw_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def coatnet_0_224(pretrained=False, **kwargs):
+def coatnet_0_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('coatnet_0_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def coatnet_1_224(pretrained=False, **kwargs):
+def coatnet_1_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('coatnet_1_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def coatnet_2_224(pretrained=False, **kwargs):
+def coatnet_2_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('coatnet_2_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def coatnet_3_224(pretrained=False, **kwargs):
+def coatnet_3_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('coatnet_3_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def coatnet_4_224(pretrained=False, **kwargs):
+def coatnet_4_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('coatnet_4_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def coatnet_5_224(pretrained=False, **kwargs):
+def coatnet_5_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('coatnet_5_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_pico_rw_256(pretrained=False, **kwargs):
+def maxvit_pico_rw_256(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_pico_rw_256', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_nano_rw_256(pretrained=False, **kwargs):
+def maxvit_nano_rw_256(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_nano_rw_256', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_tiny_rw_224(pretrained=False, **kwargs):
+def maxvit_tiny_rw_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_tiny_rw_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_tiny_rw_256(pretrained=False, **kwargs):
+def maxvit_tiny_rw_256(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_tiny_rw_256', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_rmlp_pico_rw_256(pretrained=False, **kwargs):
+def maxvit_rmlp_pico_rw_256(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_rmlp_pico_rw_256', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_rmlp_nano_rw_256(pretrained=False, **kwargs):
+def maxvit_rmlp_nano_rw_256(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_rmlp_nano_rw_256', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_rmlp_tiny_rw_256(pretrained=False, **kwargs):
+def maxvit_rmlp_tiny_rw_256(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_rmlp_tiny_rw_256', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_rmlp_small_rw_224(pretrained=False, **kwargs):
+def maxvit_rmlp_small_rw_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_rmlp_small_rw_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_rmlp_small_rw_256(pretrained=False, **kwargs):
+def maxvit_rmlp_small_rw_256(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_rmlp_small_rw_256', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_rmlp_base_rw_224(pretrained=False, **kwargs):
+def maxvit_rmlp_base_rw_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_rmlp_base_rw_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_rmlp_base_rw_384(pretrained=False, **kwargs):
+def maxvit_rmlp_base_rw_384(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_rmlp_base_rw_384', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_tiny_pm_256(pretrained=False, **kwargs):
+def maxvit_tiny_pm_256(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_tiny_pm_256', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxxvit_rmlp_nano_rw_256(pretrained=False, **kwargs):
+def maxxvit_rmlp_nano_rw_256(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxxvit_rmlp_nano_rw_256', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxxvit_rmlp_tiny_rw_256(pretrained=False, **kwargs):
+def maxxvit_rmlp_tiny_rw_256(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxxvit_rmlp_tiny_rw_256', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxxvit_rmlp_small_rw_256(pretrained=False, **kwargs):
+def maxxvit_rmlp_small_rw_256(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxxvit_rmlp_small_rw_256', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxxvitv2_nano_rw_256(pretrained=False, **kwargs):
+def maxxvitv2_nano_rw_256(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxxvitv2_nano_rw_256', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxxvitv2_rmlp_base_rw_224(pretrained=False, **kwargs):
+def maxxvitv2_rmlp_base_rw_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxxvitv2_rmlp_base_rw_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxxvitv2_rmlp_base_rw_384(pretrained=False, **kwargs):
+def maxxvitv2_rmlp_base_rw_384(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxxvitv2_rmlp_base_rw_384', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxxvitv2_rmlp_large_rw_224(pretrained=False, **kwargs):
+def maxxvitv2_rmlp_large_rw_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxxvitv2_rmlp_large_rw_224', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_tiny_tf_224(pretrained=False, **kwargs):
+def maxvit_tiny_tf_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_tiny_tf_224', 'maxvit_tiny_tf', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_tiny_tf_384(pretrained=False, **kwargs):
+def maxvit_tiny_tf_384(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_tiny_tf_384', 'maxvit_tiny_tf', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_tiny_tf_512(pretrained=False, **kwargs):
+def maxvit_tiny_tf_512(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_tiny_tf_512', 'maxvit_tiny_tf', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_small_tf_224(pretrained=False, **kwargs):
+def maxvit_small_tf_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_small_tf_224', 'maxvit_small_tf', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_small_tf_384(pretrained=False, **kwargs):
+def maxvit_small_tf_384(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_small_tf_384', 'maxvit_small_tf', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_small_tf_512(pretrained=False, **kwargs):
+def maxvit_small_tf_512(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_small_tf_512', 'maxvit_small_tf', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_base_tf_224(pretrained=False, **kwargs):
+def maxvit_base_tf_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_base_tf_224', 'maxvit_base_tf', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_base_tf_384(pretrained=False, **kwargs):
+def maxvit_base_tf_384(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_base_tf_384', 'maxvit_base_tf', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_base_tf_512(pretrained=False, **kwargs):
+def maxvit_base_tf_512(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_base_tf_512', 'maxvit_base_tf', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_large_tf_224(pretrained=False, **kwargs):
+def maxvit_large_tf_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_large_tf_224', 'maxvit_large_tf', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_large_tf_384(pretrained=False, **kwargs):
+def maxvit_large_tf_384(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_large_tf_384', 'maxvit_large_tf', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_large_tf_512(pretrained=False, **kwargs):
+def maxvit_large_tf_512(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_large_tf_512', 'maxvit_large_tf', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_xlarge_tf_224(pretrained=False, **kwargs):
+def maxvit_xlarge_tf_224(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_xlarge_tf_224', 'maxvit_xlarge_tf', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_xlarge_tf_384(pretrained=False, **kwargs):
+def maxvit_xlarge_tf_384(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_xlarge_tf_384', 'maxvit_xlarge_tf', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def maxvit_xlarge_tf_512(pretrained=False, **kwargs):
+def maxvit_xlarge_tf_512(pretrained=False, **kwargs) -> MaxxVit:
     return _create_maxxvit('maxvit_xlarge_tf_512', 'maxvit_xlarge_tf', pretrained=pretrained, **kwargs)
diff --git a/timm/models/metaformer.py b/timm/models/metaformer.py
index 72ff365f58..98a79f598b 100644
--- a/timm/models/metaformer.py
+++ b/timm/models/metaformer.py
@@ -834,7 +834,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def poolformer_s12(pretrained=False, **kwargs):
+def poolformer_s12(pretrained=False, **kwargs) -> MetaFormer:
     model_kwargs = dict(
         depths=[2, 2, 6, 2],
         dims=[64, 128, 320, 512],
@@ -850,7 +850,7 @@ def poolformer_s12(pretrained=False, **kwargs):
 
 
 @register_model
-def poolformer_s24(pretrained=False, **kwargs):
+def poolformer_s24(pretrained=False, **kwargs) -> MetaFormer:
     model_kwargs = dict(
         depths=[4, 4, 12, 4],
         dims=[64, 128, 320, 512],
@@ -866,7 +866,7 @@ def poolformer_s24(pretrained=False, **kwargs):
 
 
 @register_model
-def poolformer_s36(pretrained=False, **kwargs):
+def poolformer_s36(pretrained=False, **kwargs) -> MetaFormer:
     model_kwargs = dict(
         depths=[6, 6, 18, 6],
         dims=[64, 128, 320, 512],
@@ -882,7 +882,7 @@ def poolformer_s36(pretrained=False, **kwargs):
 
 
 @register_model
-def poolformer_m36(pretrained=False, **kwargs):
+def poolformer_m36(pretrained=False, **kwargs) -> MetaFormer:
     model_kwargs = dict(
         depths=[6, 6, 18, 6],
         dims=[96, 192, 384, 768],
@@ -898,7 +898,7 @@ def poolformer_m36(pretrained=False, **kwargs):
 
 
 @register_model
-def poolformer_m48(pretrained=False, **kwargs):
+def poolformer_m48(pretrained=False, **kwargs) -> MetaFormer:
     model_kwargs = dict(
         depths=[8, 8, 24, 8],
         dims=[96, 192, 384, 768],
@@ -914,7 +914,7 @@ def poolformer_m48(pretrained=False, **kwargs):
 
 
 @register_model
-def poolformerv2_s12(pretrained=False, **kwargs):
+def poolformerv2_s12(pretrained=False, **kwargs) -> MetaFormer:
     model_kwargs = dict(
         depths=[2, 2, 6, 2],
         dims=[64, 128, 320, 512],
@@ -925,7 +925,7 @@ def poolformerv2_s12(pretrained=False, **kwargs):
 
 
 @register_model
-def poolformerv2_s24(pretrained=False, **kwargs):
+def poolformerv2_s24(pretrained=False, **kwargs) -> MetaFormer:
     model_kwargs = dict(
         depths=[4, 4, 12, 4],
         dims=[64, 128, 320, 512],
@@ -936,7 +936,7 @@ def poolformerv2_s24(pretrained=False, **kwargs):
 
 
 @register_model
-def poolformerv2_s36(pretrained=False, **kwargs):
+def poolformerv2_s36(pretrained=False, **kwargs) -> MetaFormer:
     model_kwargs = dict(
         depths=[6, 6, 18, 6],
         dims=[64, 128, 320, 512],
@@ -947,7 +947,7 @@ def poolformerv2_s36(pretrained=False, **kwargs):
 
 
 @register_model
-def poolformerv2_m36(pretrained=False, **kwargs):
+def poolformerv2_m36(pretrained=False, **kwargs) -> MetaFormer:
     model_kwargs = dict(
         depths=[6, 6, 18, 6],
         dims=[96, 192, 384, 768],
@@ -958,7 +958,7 @@ def poolformerv2_m36(pretrained=False, **kwargs):
 
 
 @register_model
-def poolformerv2_m48(pretrained=False, **kwargs):
+def poolformerv2_m48(pretrained=False, **kwargs) -> MetaFormer:
     model_kwargs = dict(
         depths=[8, 8, 24, 8],
         dims=[96, 192, 384, 768],
@@ -969,7 +969,7 @@ def poolformerv2_m48(pretrained=False, **kwargs):
 
 
 @register_model
-def convformer_s18(pretrained=False, **kwargs):
+def convformer_s18(pretrained=False, **kwargs) -> MetaFormer:
     model_kwargs = dict(
         depths=[3, 3, 9, 3],
         dims=[64, 128, 320, 512],
@@ -980,7 +980,7 @@ def convformer_s18(pretrained=False, **kwargs):
 
 
 @register_model
-def convformer_s36(pretrained=False, **kwargs):
+def convformer_s36(pretrained=False, **kwargs) -> MetaFormer:
     model_kwargs = dict(
         depths=[3, 12, 18, 3],
         dims=[64, 128, 320, 512],
@@ -991,7 +991,7 @@ def convformer_s36(pretrained=False, **kwargs):
 
 
 @register_model
-def convformer_m36(pretrained=False, **kwargs):
+def convformer_m36(pretrained=False, **kwargs) -> MetaFormer:
     model_kwargs = dict(
         depths=[3, 12, 18, 3],
         dims=[96, 192, 384, 576],
@@ -1002,7 +1002,7 @@ def convformer_m36(pretrained=False, **kwargs):
 
 
 @register_model
-def convformer_b36(pretrained=False, **kwargs):
+def convformer_b36(pretrained=False, **kwargs) -> MetaFormer:
     model_kwargs = dict(
         depths=[3, 12, 18, 3],
         dims=[128, 256, 512, 768],
@@ -1013,7 +1013,7 @@ def convformer_b36(pretrained=False, **kwargs):
 
 
 @register_model
-def caformer_s18(pretrained=False, **kwargs):
+def caformer_s18(pretrained=False, **kwargs) -> MetaFormer:
     model_kwargs = dict(
         depths=[3, 3, 9, 3],
         dims=[64, 128, 320, 512],
@@ -1024,7 +1024,7 @@ def caformer_s18(pretrained=False, **kwargs):
 
 
 @register_model
-def caformer_s36(pretrained=False, **kwargs):
+def caformer_s36(pretrained=False, **kwargs) -> MetaFormer:
     model_kwargs = dict(
         depths=[3, 12, 18, 3],
         dims=[64, 128, 320, 512],
@@ -1035,7 +1035,7 @@ def caformer_s36(pretrained=False, **kwargs):
 
 
 @register_model
-def caformer_m36(pretrained=False, **kwargs):
+def caformer_m36(pretrained=False, **kwargs) -> MetaFormer:
     model_kwargs = dict(
         depths=[3, 12, 18, 3],
         dims=[96, 192, 384, 576],
@@ -1046,7 +1046,7 @@ def caformer_m36(pretrained=False, **kwargs):
 
 
 @register_model
-def caformer_b36(pretrained=False, **kwargs):
+def caformer_b36(pretrained=False, **kwargs) -> MetaFormer:
     model_kwargs = dict(
         depths=[3, 12, 18, 3],
         dims=[128, 256, 512, 768],
diff --git a/timm/models/mlp_mixer.py b/timm/models/mlp_mixer.py
index 2f8d55faf8..f7d64349d4 100644
--- a/timm/models/mlp_mixer.py
+++ b/timm/models/mlp_mixer.py
@@ -455,7 +455,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def mixer_s32_224(pretrained=False, **kwargs):
+def mixer_s32_224(pretrained=False, **kwargs) -> MlpMixer:
     """ Mixer-S/32 224x224
     Paper: 'MLP-Mixer: An all-MLP Architecture for Vision' - https://arxiv.org/abs/2105.01601
     """
@@ -465,7 +465,7 @@ def mixer_s32_224(pretrained=False, **kwargs):
 
 
 @register_model
-def mixer_s16_224(pretrained=False, **kwargs):
+def mixer_s16_224(pretrained=False, **kwargs) -> MlpMixer:
     """ Mixer-S/16 224x224
     Paper:  'MLP-Mixer: An all-MLP Architecture for Vision' - https://arxiv.org/abs/2105.01601
     """
@@ -475,7 +475,7 @@ def mixer_s16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def mixer_b32_224(pretrained=False, **kwargs):
+def mixer_b32_224(pretrained=False, **kwargs) -> MlpMixer:
     """ Mixer-B/32 224x224
     Paper:  'MLP-Mixer: An all-MLP Architecture for Vision' - https://arxiv.org/abs/2105.01601
     """
@@ -485,7 +485,7 @@ def mixer_b32_224(pretrained=False, **kwargs):
 
 
 @register_model
-def mixer_b16_224(pretrained=False, **kwargs):
+def mixer_b16_224(pretrained=False, **kwargs) -> MlpMixer:
     """ Mixer-B/16 224x224. ImageNet-1k pretrained weights.
     Paper:  'MLP-Mixer: An all-MLP Architecture for Vision' - https://arxiv.org/abs/2105.01601
     """
@@ -495,7 +495,7 @@ def mixer_b16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def mixer_l32_224(pretrained=False, **kwargs):
+def mixer_l32_224(pretrained=False, **kwargs) -> MlpMixer:
     """ Mixer-L/32 224x224.
     Paper:  'MLP-Mixer: An all-MLP Architecture for Vision' - https://arxiv.org/abs/2105.01601
     """
@@ -505,7 +505,7 @@ def mixer_l32_224(pretrained=False, **kwargs):
 
 
 @register_model
-def mixer_l16_224(pretrained=False, **kwargs):
+def mixer_l16_224(pretrained=False, **kwargs) -> MlpMixer:
     """ Mixer-L/16 224x224. ImageNet-1k pretrained weights.
     Paper:  'MLP-Mixer: An all-MLP Architecture for Vision' - https://arxiv.org/abs/2105.01601
     """
@@ -515,7 +515,7 @@ def mixer_l16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def gmixer_12_224(pretrained=False, **kwargs):
+def gmixer_12_224(pretrained=False, **kwargs) -> MlpMixer:
     """ Glu-Mixer-12 224x224
     Experiment by Ross Wightman, adding SwiGLU to MLP-Mixer
     """
@@ -527,7 +527,7 @@ def gmixer_12_224(pretrained=False, **kwargs):
 
 
 @register_model
-def gmixer_24_224(pretrained=False, **kwargs):
+def gmixer_24_224(pretrained=False, **kwargs) -> MlpMixer:
     """ Glu-Mixer-24 224x224
     Experiment by Ross Wightman, adding SwiGLU to MLP-Mixer
     """
@@ -539,7 +539,7 @@ def gmixer_24_224(pretrained=False, **kwargs):
 
 
 @register_model
-def resmlp_12_224(pretrained=False, **kwargs):
+def resmlp_12_224(pretrained=False, **kwargs) -> MlpMixer:
     """ ResMLP-12
     Paper: `ResMLP: Feedforward networks for image classification...` - https://arxiv.org/abs/2105.03404
     """
@@ -550,7 +550,7 @@ def resmlp_12_224(pretrained=False, **kwargs):
 
 
 @register_model
-def resmlp_24_224(pretrained=False, **kwargs):
+def resmlp_24_224(pretrained=False, **kwargs) -> MlpMixer:
     """ ResMLP-24
     Paper: `ResMLP: Feedforward networks for image classification...` - https://arxiv.org/abs/2105.03404
     """
@@ -562,7 +562,7 @@ def resmlp_24_224(pretrained=False, **kwargs):
 
 
 @register_model
-def resmlp_36_224(pretrained=False, **kwargs):
+def resmlp_36_224(pretrained=False, **kwargs) -> MlpMixer:
     """ ResMLP-36
     Paper: `ResMLP: Feedforward networks for image classification...` - https://arxiv.org/abs/2105.03404
     """
@@ -574,7 +574,7 @@ def resmlp_36_224(pretrained=False, **kwargs):
 
 
 @register_model
-def resmlp_big_24_224(pretrained=False, **kwargs):
+def resmlp_big_24_224(pretrained=False, **kwargs) -> MlpMixer:
     """ ResMLP-B-24
     Paper: `ResMLP: Feedforward networks for image classification...` - https://arxiv.org/abs/2105.03404
     """
@@ -586,7 +586,7 @@ def resmlp_big_24_224(pretrained=False, **kwargs):
 
 
 @register_model
-def gmlp_ti16_224(pretrained=False, **kwargs):
+def gmlp_ti16_224(pretrained=False, **kwargs) -> MlpMixer:
     """ gMLP-Tiny
     Paper: `Pay Attention to MLPs` - https://arxiv.org/abs/2105.08050
     """
@@ -598,7 +598,7 @@ def gmlp_ti16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def gmlp_s16_224(pretrained=False, **kwargs):
+def gmlp_s16_224(pretrained=False, **kwargs) -> MlpMixer:
     """ gMLP-Small
     Paper: `Pay Attention to MLPs` - https://arxiv.org/abs/2105.08050
     """
@@ -610,7 +610,7 @@ def gmlp_s16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def gmlp_b16_224(pretrained=False, **kwargs):
+def gmlp_b16_224(pretrained=False, **kwargs) -> MlpMixer:
     """ gMLP-Base
     Paper: `Pay Attention to MLPs` - https://arxiv.org/abs/2105.08050
     """
diff --git a/timm/models/mobilenetv3.py b/timm/models/mobilenetv3.py
index 849d0c21d8..70087b62fa 100644
--- a/timm/models/mobilenetv3.py
+++ b/timm/models/mobilenetv3.py
@@ -643,42 +643,42 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def mobilenetv3_large_075(pretrained=False, **kwargs):
+def mobilenetv3_large_075(pretrained=False, **kwargs) -> MobileNetV3:
     """ MobileNet V3 """
     model = _gen_mobilenet_v3('mobilenetv3_large_075', 0.75, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def mobilenetv3_large_100(pretrained=False, **kwargs):
+def mobilenetv3_large_100(pretrained=False, **kwargs) -> MobileNetV3:
     """ MobileNet V3 """
     model = _gen_mobilenet_v3('mobilenetv3_large_100', 1.0, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def mobilenetv3_small_050(pretrained=False, **kwargs):
+def mobilenetv3_small_050(pretrained=False, **kwargs) -> MobileNetV3:
     """ MobileNet V3 """
     model = _gen_mobilenet_v3('mobilenetv3_small_050', 0.50, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def mobilenetv3_small_075(pretrained=False, **kwargs):
+def mobilenetv3_small_075(pretrained=False, **kwargs) -> MobileNetV3:
     """ MobileNet V3 """
     model = _gen_mobilenet_v3('mobilenetv3_small_075', 0.75, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def mobilenetv3_small_100(pretrained=False, **kwargs):
+def mobilenetv3_small_100(pretrained=False, **kwargs) -> MobileNetV3:
     """ MobileNet V3 """
     model = _gen_mobilenet_v3('mobilenetv3_small_100', 1.0, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def mobilenetv3_rw(pretrained=False, **kwargs):
+def mobilenetv3_rw(pretrained=False, **kwargs) -> MobileNetV3:
     """ MobileNet V3 """
     if pretrained:
         # pretrained model trained with non-default BN epsilon
@@ -688,7 +688,7 @@ def mobilenetv3_rw(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_mobilenetv3_large_075(pretrained=False, **kwargs):
+def tf_mobilenetv3_large_075(pretrained=False, **kwargs) -> MobileNetV3:
     """ MobileNet V3 """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -697,7 +697,7 @@ def tf_mobilenetv3_large_075(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_mobilenetv3_large_100(pretrained=False, **kwargs):
+def tf_mobilenetv3_large_100(pretrained=False, **kwargs) -> MobileNetV3:
     """ MobileNet V3 """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -706,7 +706,7 @@ def tf_mobilenetv3_large_100(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_mobilenetv3_large_minimal_100(pretrained=False, **kwargs):
+def tf_mobilenetv3_large_minimal_100(pretrained=False, **kwargs) -> MobileNetV3:
     """ MobileNet V3 """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -715,7 +715,7 @@ def tf_mobilenetv3_large_minimal_100(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_mobilenetv3_small_075(pretrained=False, **kwargs):
+def tf_mobilenetv3_small_075(pretrained=False, **kwargs) -> MobileNetV3:
     """ MobileNet V3 """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -724,7 +724,7 @@ def tf_mobilenetv3_small_075(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_mobilenetv3_small_100(pretrained=False, **kwargs):
+def tf_mobilenetv3_small_100(pretrained=False, **kwargs) -> MobileNetV3:
     """ MobileNet V3 """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -733,7 +733,7 @@ def tf_mobilenetv3_small_100(pretrained=False, **kwargs):
 
 
 @register_model
-def tf_mobilenetv3_small_minimal_100(pretrained=False, **kwargs):
+def tf_mobilenetv3_small_minimal_100(pretrained=False, **kwargs) -> MobileNetV3:
     """ MobileNet V3 """
     kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
     kwargs['pad_type'] = 'same'
@@ -742,56 +742,56 @@ def tf_mobilenetv3_small_minimal_100(pretrained=False, **kwargs):
 
 
 @register_model
-def fbnetv3_b(pretrained=False, **kwargs):
+def fbnetv3_b(pretrained=False, **kwargs) -> MobileNetV3:
     """ FBNetV3-B """
     model = _gen_fbnetv3('fbnetv3_b', pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def fbnetv3_d(pretrained=False, **kwargs):
+def fbnetv3_d(pretrained=False, **kwargs) -> MobileNetV3:
     """ FBNetV3-D """
     model = _gen_fbnetv3('fbnetv3_d', pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def fbnetv3_g(pretrained=False, **kwargs):
+def fbnetv3_g(pretrained=False, **kwargs) -> MobileNetV3:
     """ FBNetV3-G """
     model = _gen_fbnetv3('fbnetv3_g', pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def lcnet_035(pretrained=False, **kwargs):
+def lcnet_035(pretrained=False, **kwargs) -> MobileNetV3:
     """ PP-LCNet 0.35"""
     model = _gen_lcnet('lcnet_035', 0.35, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def lcnet_050(pretrained=False, **kwargs):
+def lcnet_050(pretrained=False, **kwargs) -> MobileNetV3:
     """ PP-LCNet 0.5"""
     model = _gen_lcnet('lcnet_050', 0.5, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def lcnet_075(pretrained=False, **kwargs):
+def lcnet_075(pretrained=False, **kwargs) -> MobileNetV3:
     """ PP-LCNet 1.0"""
     model = _gen_lcnet('lcnet_075', 0.75, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def lcnet_100(pretrained=False, **kwargs):
+def lcnet_100(pretrained=False, **kwargs) -> MobileNetV3:
     """ PP-LCNet 1.0"""
     model = _gen_lcnet('lcnet_100', 1.0, pretrained=pretrained, **kwargs)
     return model
 
 
 @register_model
-def lcnet_150(pretrained=False, **kwargs):
+def lcnet_150(pretrained=False, **kwargs) -> MobileNetV3:
     """ PP-LCNet 1.5"""
     model = _gen_lcnet('lcnet_150', 1.5, pretrained=pretrained, **kwargs)
     return model
diff --git a/timm/models/mobilevit.py b/timm/models/mobilevit.py
index 9fb986a6e4..9c84871e6d 100644
--- a/timm/models/mobilevit.py
+++ b/timm/models/mobilevit.py
@@ -621,52 +621,52 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def mobilevit_xxs(pretrained=False, **kwargs):
+def mobilevit_xxs(pretrained=False, **kwargs) -> ByobNet:
     return _create_mobilevit('mobilevit_xxs', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def mobilevit_xs(pretrained=False, **kwargs):
+def mobilevit_xs(pretrained=False, **kwargs) -> ByobNet:
     return _create_mobilevit('mobilevit_xs', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def mobilevit_s(pretrained=False, **kwargs):
+def mobilevit_s(pretrained=False, **kwargs) -> ByobNet:
     return _create_mobilevit('mobilevit_s', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def mobilevitv2_050(pretrained=False, **kwargs):
+def mobilevitv2_050(pretrained=False, **kwargs) -> ByobNet:
     return _create_mobilevit('mobilevitv2_050', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def mobilevitv2_075(pretrained=False, **kwargs):
+def mobilevitv2_075(pretrained=False, **kwargs) -> ByobNet:
     return _create_mobilevit('mobilevitv2_075', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def mobilevitv2_100(pretrained=False, **kwargs):
+def mobilevitv2_100(pretrained=False, **kwargs) -> ByobNet:
     return _create_mobilevit('mobilevitv2_100', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def mobilevitv2_125(pretrained=False, **kwargs):
+def mobilevitv2_125(pretrained=False, **kwargs) -> ByobNet:
     return _create_mobilevit('mobilevitv2_125', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def mobilevitv2_150(pretrained=False, **kwargs):
+def mobilevitv2_150(pretrained=False, **kwargs) -> ByobNet:
     return _create_mobilevit('mobilevitv2_150', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def mobilevitv2_175(pretrained=False, **kwargs):
+def mobilevitv2_175(pretrained=False, **kwargs) -> ByobNet:
     return _create_mobilevit('mobilevitv2_175', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def mobilevitv2_200(pretrained=False, **kwargs):
+def mobilevitv2_200(pretrained=False, **kwargs) -> ByobNet:
     return _create_mobilevit('mobilevitv2_200', pretrained=pretrained, **kwargs)
 
 
diff --git a/timm/models/mvitv2.py b/timm/models/mvitv2.py
index 23d7a94721..bc18bbc2cd 100644
--- a/timm/models/mvitv2.py
+++ b/timm/models/mvitv2.py
@@ -991,40 +991,40 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def mvitv2_tiny(pretrained=False, **kwargs):
+def mvitv2_tiny(pretrained=False, **kwargs) -> MultiScaleVit:
     return _create_mvitv2('mvitv2_tiny', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def mvitv2_small(pretrained=False, **kwargs):
+def mvitv2_small(pretrained=False, **kwargs) -> MultiScaleVit:
     return _create_mvitv2('mvitv2_small', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def mvitv2_base(pretrained=False, **kwargs):
+def mvitv2_base(pretrained=False, **kwargs) -> MultiScaleVit:
     return _create_mvitv2('mvitv2_base', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def mvitv2_large(pretrained=False, **kwargs):
+def mvitv2_large(pretrained=False, **kwargs) -> MultiScaleVit:
     return _create_mvitv2('mvitv2_large', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def mvitv2_small_cls(pretrained=False, **kwargs):
+def mvitv2_small_cls(pretrained=False, **kwargs) -> MultiScaleVit:
     return _create_mvitv2('mvitv2_small_cls', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def mvitv2_base_cls(pretrained=False, **kwargs):
+def mvitv2_base_cls(pretrained=False, **kwargs) -> MultiScaleVit:
     return _create_mvitv2('mvitv2_base_cls', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def mvitv2_large_cls(pretrained=False, **kwargs):
+def mvitv2_large_cls(pretrained=False, **kwargs) -> MultiScaleVit:
     return _create_mvitv2('mvitv2_large_cls', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def mvitv2_huge_cls(pretrained=False, **kwargs):
+def mvitv2_huge_cls(pretrained=False, **kwargs) -> MultiScaleVit:
     return _create_mvitv2('mvitv2_huge_cls', pretrained=pretrained, **kwargs)
diff --git a/timm/models/nasnet.py b/timm/models/nasnet.py
index 61fdc7f752..edf355f8e0 100644
--- a/timm/models/nasnet.py
+++ b/timm/models/nasnet.py
@@ -594,7 +594,7 @@ def _create_nasnet(variant, pretrained=False, **kwargs):
 
 
 @register_model
-def nasnetalarge(pretrained=False, **kwargs):
+def nasnetalarge(pretrained=False, **kwargs) -> NASNetALarge:
     """NASNet-A large model architecture.
     """
     model_kwargs = dict(pad_type='same', **kwargs)
diff --git a/timm/models/nest.py b/timm/models/nest.py
index bc6984e69d..de57ec6e99 100644
--- a/timm/models/nest.py
+++ b/timm/models/nest.py
@@ -518,7 +518,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def nest_base(pretrained=False, **kwargs):
+def nest_base(pretrained=False, **kwargs) -> Nest:
     """ Nest-B @ 224x224
     """
     model_kwargs = dict(
@@ -528,7 +528,7 @@ def nest_base(pretrained=False, **kwargs):
 
 
 @register_model
-def nest_small(pretrained=False, **kwargs):
+def nest_small(pretrained=False, **kwargs) -> Nest:
     """ Nest-S @ 224x224
     """
     model_kwargs = dict(embed_dims=(96, 192, 384), num_heads=(3, 6, 12), depths=(2, 2, 20), **kwargs)
@@ -537,7 +537,7 @@ def nest_small(pretrained=False, **kwargs):
 
 
 @register_model
-def nest_tiny(pretrained=False, **kwargs):
+def nest_tiny(pretrained=False, **kwargs) -> Nest:
     """ Nest-T @ 224x224
     """
     model_kwargs = dict(embed_dims=(96, 192, 384), num_heads=(3, 6, 12), depths=(2, 2, 8), **kwargs)
@@ -546,7 +546,7 @@ def nest_tiny(pretrained=False, **kwargs):
 
 
 @register_model
-def nest_base_jx(pretrained=False, **kwargs):
+def nest_base_jx(pretrained=False, **kwargs) -> Nest:
     """ Nest-B @ 224x224
     """
     kwargs.setdefault('pad_type', 'same')
@@ -557,7 +557,7 @@ def nest_base_jx(pretrained=False, **kwargs):
 
 
 @register_model
-def nest_small_jx(pretrained=False, **kwargs):
+def nest_small_jx(pretrained=False, **kwargs) -> Nest:
     """ Nest-S @ 224x224
     """
     kwargs.setdefault('pad_type', 'same')
@@ -567,7 +567,7 @@ def nest_small_jx(pretrained=False, **kwargs):
 
 
 @register_model
-def nest_tiny_jx(pretrained=False, **kwargs):
+def nest_tiny_jx(pretrained=False, **kwargs) -> Nest:
     """ Nest-T @ 224x224
     """
     kwargs.setdefault('pad_type', 'same')
diff --git a/timm/models/nfnet.py b/timm/models/nfnet.py
index d46e02794a..725b177c25 100644
--- a/timm/models/nfnet.py
+++ b/timm/models/nfnet.py
@@ -733,7 +733,7 @@ def _dcfg(url='', **kwargs):
 
 
 @register_model
-def dm_nfnet_f0(pretrained=False, **kwargs):
+def dm_nfnet_f0(pretrained=False, **kwargs) -> NormFreeNet:
     """ NFNet-F0 (DeepMind weight compatible)
     `High-Performance Large-Scale Image Recognition Without Normalization`
         - https://arxiv.org/abs/2102.06171
@@ -742,7 +742,7 @@ def dm_nfnet_f0(pretrained=False, **kwargs):
 
 
 @register_model
-def dm_nfnet_f1(pretrained=False, **kwargs):
+def dm_nfnet_f1(pretrained=False, **kwargs) -> NormFreeNet:
     """ NFNet-F1 (DeepMind weight compatible)
     `High-Performance Large-Scale Image Recognition Without Normalization`
         - https://arxiv.org/abs/2102.06171
@@ -751,7 +751,7 @@ def dm_nfnet_f1(pretrained=False, **kwargs):
 
 
 @register_model
-def dm_nfnet_f2(pretrained=False, **kwargs):
+def dm_nfnet_f2(pretrained=False, **kwargs) -> NormFreeNet:
     """ NFNet-F2 (DeepMind weight compatible)
     `High-Performance Large-Scale Image Recognition Without Normalization`
         - https://arxiv.org/abs/2102.06171
@@ -760,7 +760,7 @@ def dm_nfnet_f2(pretrained=False, **kwargs):
 
 
 @register_model
-def dm_nfnet_f3(pretrained=False, **kwargs):
+def dm_nfnet_f3(pretrained=False, **kwargs) -> NormFreeNet:
     """ NFNet-F3 (DeepMind weight compatible)
     `High-Performance Large-Scale Image Recognition Without Normalization`
         - https://arxiv.org/abs/2102.06171
@@ -769,7 +769,7 @@ def dm_nfnet_f3(pretrained=False, **kwargs):
 
 
 @register_model
-def dm_nfnet_f4(pretrained=False, **kwargs):
+def dm_nfnet_f4(pretrained=False, **kwargs) -> NormFreeNet:
     """ NFNet-F4 (DeepMind weight compatible)
     `High-Performance Large-Scale Image Recognition Without Normalization`
         - https://arxiv.org/abs/2102.06171
@@ -778,7 +778,7 @@ def dm_nfnet_f4(pretrained=False, **kwargs):
 
 
 @register_model
-def dm_nfnet_f5(pretrained=False, **kwargs):
+def dm_nfnet_f5(pretrained=False, **kwargs) -> NormFreeNet:
     """ NFNet-F5 (DeepMind weight compatible)
     `High-Performance Large-Scale Image Recognition Without Normalization`
         - https://arxiv.org/abs/2102.06171
@@ -787,7 +787,7 @@ def dm_nfnet_f5(pretrained=False, **kwargs):
 
 
 @register_model
-def dm_nfnet_f6(pretrained=False, **kwargs):
+def dm_nfnet_f6(pretrained=False, **kwargs) -> NormFreeNet:
     """ NFNet-F6 (DeepMind weight compatible)
     `High-Performance Large-Scale Image Recognition Without Normalization`
         - https://arxiv.org/abs/2102.06171
@@ -796,7 +796,7 @@ def dm_nfnet_f6(pretrained=False, **kwargs):
 
 
 @register_model
-def nfnet_f0(pretrained=False, **kwargs):
+def nfnet_f0(pretrained=False, **kwargs) -> NormFreeNet:
     """ NFNet-F0
     `High-Performance Large-Scale Image Recognition Without Normalization`
         - https://arxiv.org/abs/2102.06171
@@ -805,7 +805,7 @@ def nfnet_f0(pretrained=False, **kwargs):
 
 
 @register_model
-def nfnet_f1(pretrained=False, **kwargs):
+def nfnet_f1(pretrained=False, **kwargs) -> NormFreeNet:
     """ NFNet-F1
     `High-Performance Large-Scale Image Recognition Without Normalization`
         - https://arxiv.org/abs/2102.06171
@@ -814,7 +814,7 @@ def nfnet_f1(pretrained=False, **kwargs):
 
 
 @register_model
-def nfnet_f2(pretrained=False, **kwargs):
+def nfnet_f2(pretrained=False, **kwargs) -> NormFreeNet:
     """ NFNet-F2
     `High-Performance Large-Scale Image Recognition Without Normalization`
         - https://arxiv.org/abs/2102.06171
@@ -823,7 +823,7 @@ def nfnet_f2(pretrained=False, **kwargs):
 
 
 @register_model
-def nfnet_f3(pretrained=False, **kwargs):
+def nfnet_f3(pretrained=False, **kwargs) -> NormFreeNet:
     """ NFNet-F3
     `High-Performance Large-Scale Image Recognition Without Normalization`
         - https://arxiv.org/abs/2102.06171
@@ -832,7 +832,7 @@ def nfnet_f3(pretrained=False, **kwargs):
 
 
 @register_model
-def nfnet_f4(pretrained=False, **kwargs):
+def nfnet_f4(pretrained=False, **kwargs) -> NormFreeNet:
     """ NFNet-F4
     `High-Performance Large-Scale Image Recognition Without Normalization`
         - https://arxiv.org/abs/2102.06171
@@ -841,7 +841,7 @@ def nfnet_f4(pretrained=False, **kwargs):
 
 
 @register_model
-def nfnet_f5(pretrained=False, **kwargs):
+def nfnet_f5(pretrained=False, **kwargs) -> NormFreeNet:
     """ NFNet-F5
     `High-Performance Large-Scale Image Recognition Without Normalization`
         - https://arxiv.org/abs/2102.06171
@@ -850,7 +850,7 @@ def nfnet_f5(pretrained=False, **kwargs):
 
 
 @register_model
-def nfnet_f6(pretrained=False, **kwargs):
+def nfnet_f6(pretrained=False, **kwargs) -> NormFreeNet:
     """ NFNet-F6
     `High-Performance Large-Scale Image Recognition Without Normalization`
         - https://arxiv.org/abs/2102.06171
@@ -859,7 +859,7 @@ def nfnet_f6(pretrained=False, **kwargs):
 
 
 @register_model
-def nfnet_f7(pretrained=False, **kwargs):
+def nfnet_f7(pretrained=False, **kwargs) -> NormFreeNet:
     """ NFNet-F7
     `High-Performance Large-Scale Image Recognition Without Normalization`
         - https://arxiv.org/abs/2102.06171
@@ -868,7 +868,7 @@ def nfnet_f7(pretrained=False, **kwargs):
 
 
 @register_model
-def nfnet_l0(pretrained=False, **kwargs):
+def nfnet_l0(pretrained=False, **kwargs) -> NormFreeNet:
     """ NFNet-L0b w/ SiLU
     My experimental 'light' model w/ F0 repeats, 1.5x final_conv mult, 64 group_size, .25 bottleneck & SE ratio
     """
@@ -876,7 +876,7 @@ def nfnet_l0(pretrained=False, **kwargs):
 
 
 @register_model
-def eca_nfnet_l0(pretrained=False, **kwargs):
+def eca_nfnet_l0(pretrained=False, **kwargs) -> NormFreeNet:
     """ ECA-NFNet-L0 w/ SiLU
     My experimental 'light' model w/ F0 repeats, 1.5x final_conv mult, 64 group_size, .25 bottleneck & ECA attn
     """
@@ -884,7 +884,7 @@ def eca_nfnet_l0(pretrained=False, **kwargs):
 
 
 @register_model
-def eca_nfnet_l1(pretrained=False, **kwargs):
+def eca_nfnet_l1(pretrained=False, **kwargs) -> NormFreeNet:
     """ ECA-NFNet-L1 w/ SiLU
     My experimental 'light' model w/ F1 repeats, 2.0x final_conv mult, 64 group_size, .25 bottleneck & ECA attn
     """
@@ -892,7 +892,7 @@ def eca_nfnet_l1(pretrained=False, **kwargs):
 
 
 @register_model
-def eca_nfnet_l2(pretrained=False, **kwargs):
+def eca_nfnet_l2(pretrained=False, **kwargs) -> NormFreeNet:
     """ ECA-NFNet-L2 w/ SiLU
     My experimental 'light' model w/ F2 repeats, 2.0x final_conv mult, 64 group_size, .25 bottleneck & ECA attn
     """
@@ -900,7 +900,7 @@ def eca_nfnet_l2(pretrained=False, **kwargs):
 
 
 @register_model
-def eca_nfnet_l3(pretrained=False, **kwargs):
+def eca_nfnet_l3(pretrained=False, **kwargs) -> NormFreeNet:
     """ ECA-NFNet-L3 w/ SiLU
     My experimental 'light' model w/ F3 repeats, 2.0x final_conv mult, 64 group_size, .25 bottleneck & ECA attn
     """
@@ -908,7 +908,7 @@ def eca_nfnet_l3(pretrained=False, **kwargs):
 
 
 @register_model
-def nf_regnet_b0(pretrained=False, **kwargs):
+def nf_regnet_b0(pretrained=False, **kwargs) -> NormFreeNet:
     """ Normalization-Free RegNet-B0
     `Characterizing signal propagation to close the performance gap in unnormalized ResNets`
         - https://arxiv.org/abs/2101.08692
@@ -917,7 +917,7 @@ def nf_regnet_b0(pretrained=False, **kwargs):
 
 
 @register_model
-def nf_regnet_b1(pretrained=False, **kwargs):
+def nf_regnet_b1(pretrained=False, **kwargs) -> NormFreeNet:
     """ Normalization-Free RegNet-B1
     `Characterizing signal propagation to close the performance gap in unnormalized ResNets`
         - https://arxiv.org/abs/2101.08692
@@ -926,7 +926,7 @@ def nf_regnet_b1(pretrained=False, **kwargs):
 
 
 @register_model
-def nf_regnet_b2(pretrained=False, **kwargs):
+def nf_regnet_b2(pretrained=False, **kwargs) -> NormFreeNet:
     """ Normalization-Free RegNet-B2
     `Characterizing signal propagation to close the performance gap in unnormalized ResNets`
         - https://arxiv.org/abs/2101.08692
@@ -935,7 +935,7 @@ def nf_regnet_b2(pretrained=False, **kwargs):
 
 
 @register_model
-def nf_regnet_b3(pretrained=False, **kwargs):
+def nf_regnet_b3(pretrained=False, **kwargs) -> NormFreeNet:
     """ Normalization-Free RegNet-B3
     `Characterizing signal propagation to close the performance gap in unnormalized ResNets`
         - https://arxiv.org/abs/2101.08692
@@ -944,7 +944,7 @@ def nf_regnet_b3(pretrained=False, **kwargs):
 
 
 @register_model
-def nf_regnet_b4(pretrained=False, **kwargs):
+def nf_regnet_b4(pretrained=False, **kwargs) -> NormFreeNet:
     """ Normalization-Free RegNet-B4
     `Characterizing signal propagation to close the performance gap in unnormalized ResNets`
         - https://arxiv.org/abs/2101.08692
@@ -953,7 +953,7 @@ def nf_regnet_b4(pretrained=False, **kwargs):
 
 
 @register_model
-def nf_regnet_b5(pretrained=False, **kwargs):
+def nf_regnet_b5(pretrained=False, **kwargs) -> NormFreeNet:
     """ Normalization-Free RegNet-B5
     `Characterizing signal propagation to close the performance gap in unnormalized ResNets`
         - https://arxiv.org/abs/2101.08692
@@ -962,7 +962,7 @@ def nf_regnet_b5(pretrained=False, **kwargs):
 
 
 @register_model
-def nf_resnet26(pretrained=False, **kwargs):
+def nf_resnet26(pretrained=False, **kwargs) -> NormFreeNet:
     """ Normalization-Free ResNet-26
     `Characterizing signal propagation to close the performance gap in unnormalized ResNets`
         - https://arxiv.org/abs/2101.08692
@@ -971,7 +971,7 @@ def nf_resnet26(pretrained=False, **kwargs):
 
 
 @register_model
-def nf_resnet50(pretrained=False, **kwargs):
+def nf_resnet50(pretrained=False, **kwargs) -> NormFreeNet:
     """ Normalization-Free ResNet-50
     `Characterizing signal propagation to close the performance gap in unnormalized ResNets`
         - https://arxiv.org/abs/2101.08692
@@ -980,7 +980,7 @@ def nf_resnet50(pretrained=False, **kwargs):
 
 
 @register_model
-def nf_resnet101(pretrained=False, **kwargs):
+def nf_resnet101(pretrained=False, **kwargs) -> NormFreeNet:
     """ Normalization-Free ResNet-101
     `Characterizing signal propagation to close the performance gap in unnormalized ResNets`
         - https://arxiv.org/abs/2101.08692
@@ -989,42 +989,42 @@ def nf_resnet101(pretrained=False, **kwargs):
 
 
 @register_model
-def nf_seresnet26(pretrained=False, **kwargs):
+def nf_seresnet26(pretrained=False, **kwargs) -> NormFreeNet:
     """ Normalization-Free SE-ResNet26
     """
     return _create_normfreenet('nf_seresnet26', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def nf_seresnet50(pretrained=False, **kwargs):
+def nf_seresnet50(pretrained=False, **kwargs) -> NormFreeNet:
     """ Normalization-Free SE-ResNet50
     """
     return _create_normfreenet('nf_seresnet50', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def nf_seresnet101(pretrained=False, **kwargs):
+def nf_seresnet101(pretrained=False, **kwargs) -> NormFreeNet:
     """ Normalization-Free SE-ResNet101
     """
     return _create_normfreenet('nf_seresnet101', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def nf_ecaresnet26(pretrained=False, **kwargs):
+def nf_ecaresnet26(pretrained=False, **kwargs) -> NormFreeNet:
     """ Normalization-Free ECA-ResNet26
     """
     return _create_normfreenet('nf_ecaresnet26', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def nf_ecaresnet50(pretrained=False, **kwargs):
+def nf_ecaresnet50(pretrained=False, **kwargs) -> NormFreeNet:
     """ Normalization-Free ECA-ResNet50
     """
     return _create_normfreenet('nf_ecaresnet50', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def nf_ecaresnet101(pretrained=False, **kwargs):
+def nf_ecaresnet101(pretrained=False, **kwargs) -> NormFreeNet:
     """ Normalization-Free ECA-ResNet101
     """
     return _create_normfreenet('nf_ecaresnet101', pretrained=pretrained, **kwargs)
diff --git a/timm/models/pit.py b/timm/models/pit.py
index 4c5addd821..993606d518 100644
--- a/timm/models/pit.py
+++ b/timm/models/pit.py
@@ -351,7 +351,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def pit_b_224(pretrained, **kwargs):
+def pit_b_224(pretrained=False, **kwargs) -> PoolingVisionTransformer:
     model_args = dict(
         patch_size=14,
         stride=7,
@@ -364,7 +364,7 @@ def pit_b_224(pretrained, **kwargs):
 
 
 @register_model
-def pit_s_224(pretrained, **kwargs):
+def pit_s_224(pretrained=False, **kwargs) -> PoolingVisionTransformer:
     model_args = dict(
         patch_size=16,
         stride=8,
@@ -377,7 +377,7 @@ def pit_s_224(pretrained, **kwargs):
 
 
 @register_model
-def pit_xs_224(pretrained, **kwargs):
+def pit_xs_224(pretrained=False, **kwargs) -> PoolingVisionTransformer:
     model_args = dict(
         patch_size=16,
         stride=8,
@@ -390,7 +390,7 @@ def pit_xs_224(pretrained, **kwargs):
 
 
 @register_model
-def pit_ti_224(pretrained, **kwargs):
+def pit_ti_224(pretrained=False, **kwargs) -> PoolingVisionTransformer:
     model_args = dict(
         patch_size=16,
         stride=8,
@@ -403,7 +403,7 @@ def pit_ti_224(pretrained, **kwargs):
 
 
 @register_model
-def pit_b_distilled_224(pretrained, **kwargs):
+def pit_b_distilled_224(pretrained=False, **kwargs) -> PoolingVisionTransformer:
     model_args = dict(
         patch_size=14,
         stride=7,
@@ -417,7 +417,7 @@ def pit_b_distilled_224(pretrained, **kwargs):
 
 
 @register_model
-def pit_s_distilled_224(pretrained, **kwargs):
+def pit_s_distilled_224(pretrained=False, **kwargs) -> PoolingVisionTransformer:
     model_args = dict(
         patch_size=16,
         stride=8,
@@ -431,7 +431,7 @@ def pit_s_distilled_224(pretrained, **kwargs):
 
 
 @register_model
-def pit_xs_distilled_224(pretrained, **kwargs):
+def pit_xs_distilled_224(pretrained=False, **kwargs) -> PoolingVisionTransformer:
     model_args = dict(
         patch_size=16,
         stride=8,
@@ -445,7 +445,7 @@ def pit_xs_distilled_224(pretrained, **kwargs):
 
 
 @register_model
-def pit_ti_distilled_224(pretrained, **kwargs):
+def pit_ti_distilled_224(pretrained=False, **kwargs) -> PoolingVisionTransformer:
     model_args = dict(
         patch_size=16,
         stride=8,
diff --git a/timm/models/pnasnet.py b/timm/models/pnasnet.py
index 0b55f354d7..23785007c6 100644
--- a/timm/models/pnasnet.py
+++ b/timm/models/pnasnet.py
@@ -371,7 +371,7 @@ def _create_pnasnet(variant, pretrained=False, **kwargs):
 
 
 @register_model
-def pnasnet5large(pretrained=False, **kwargs):
+def pnasnet5large(pretrained=False, **kwargs) -> PNASNet5Large:
     r"""PNASNet-5 model architecture from the
     `"Progressive Neural Architecture Search"
     <https://arxiv.org/abs/1712.00559>`_ paper.
diff --git a/timm/models/pvt_v2.py b/timm/models/pvt_v2.py
index 0b4c54c81e..e594e6d50f 100644
--- a/timm/models/pvt_v2.py
+++ b/timm/models/pvt_v2.py
@@ -459,44 +459,44 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def pvt_v2_b0(pretrained=False, **kwargs):
+def pvt_v2_b0(pretrained=False, **kwargs) -> PyramidVisionTransformerV2:
     model_args = dict(depths=(2, 2, 2, 2), embed_dims=(32, 64, 160, 256), num_heads=(1, 2, 5, 8))
     return _create_pvt2('pvt_v2_b0', pretrained=pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
-def pvt_v2_b1(pretrained=False, **kwargs):
+def pvt_v2_b1(pretrained=False, **kwargs) -> PyramidVisionTransformerV2:
     model_args = dict(depths=(2, 2, 2, 2), embed_dims=(64, 128, 320, 512), num_heads=(1, 2, 5, 8))
     return _create_pvt2('pvt_v2_b1', pretrained=pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
-def pvt_v2_b2(pretrained=False, **kwargs):
+def pvt_v2_b2(pretrained=False, **kwargs) -> PyramidVisionTransformerV2:
     model_args = dict(depths=(3, 4, 6, 3), embed_dims=(64, 128, 320, 512), num_heads=(1, 2, 5, 8))
     return _create_pvt2('pvt_v2_b2', pretrained=pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
-def pvt_v2_b3(pretrained=False, **kwargs):
+def pvt_v2_b3(pretrained=False, **kwargs) -> PyramidVisionTransformerV2:
     model_args = dict(depths=(3, 4, 18, 3), embed_dims=(64, 128, 320, 512), num_heads=(1, 2, 5, 8))
     return _create_pvt2('pvt_v2_b3', pretrained=pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
-def pvt_v2_b4(pretrained=False, **kwargs):
+def pvt_v2_b4(pretrained=False, **kwargs) -> PyramidVisionTransformerV2:
     model_args = dict(depths=(3, 8, 27, 3), embed_dims=(64, 128, 320, 512), num_heads=(1, 2, 5, 8))
     return _create_pvt2('pvt_v2_b4', pretrained=pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
-def pvt_v2_b5(pretrained=False, **kwargs):
+def pvt_v2_b5(pretrained=False, **kwargs) -> PyramidVisionTransformerV2:
     model_args = dict(
         depths=(3, 6, 40, 3), embed_dims=(64, 128, 320, 512), num_heads=(1, 2, 5, 8), mlp_ratios=(4, 4, 4, 4))
     return _create_pvt2('pvt_v2_b5', pretrained=pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
-def pvt_v2_b2_li(pretrained=False, **kwargs):
+def pvt_v2_b2_li(pretrained=False, **kwargs) -> PyramidVisionTransformerV2:
     model_args = dict(
         depths=(3, 4, 6, 3), embed_dims=(64, 128, 320, 512), num_heads=(1, 2, 5, 8), linear=True)
     return _create_pvt2('pvt_v2_b2_li', pretrained=pretrained, **dict(model_args, **kwargs))
diff --git a/timm/models/regnet.py b/timm/models/regnet.py
index 250ba2e449..4ece9f4c01 100644
--- a/timm/models/regnet.py
+++ b/timm/models/regnet.py
@@ -893,205 +893,205 @@ def _cfgtv2(url='', **kwargs):
 
 
 @register_model
-def regnetx_002(pretrained=False, **kwargs):
+def regnetx_002(pretrained=False, **kwargs) -> RegNet:
     """RegNetX-200MF"""
     return _create_regnet('regnetx_002', pretrained, **kwargs)
 
 
 @register_model
-def regnetx_004(pretrained=False, **kwargs):
+def regnetx_004(pretrained=False, **kwargs) -> RegNet:
     """RegNetX-400MF"""
     return _create_regnet('regnetx_004', pretrained, **kwargs)
 
 
 @register_model
-def regnetx_004_tv(pretrained=False, **kwargs):
+def regnetx_004_tv(pretrained=False, **kwargs) -> RegNet:
     """RegNetX-400MF w/ torchvision group rounding"""
     return _create_regnet('regnetx_004_tv', pretrained, **kwargs)
 
 
 @register_model
-def regnetx_006(pretrained=False, **kwargs):
+def regnetx_006(pretrained=False, **kwargs) -> RegNet:
     """RegNetX-600MF"""
     return _create_regnet('regnetx_006', pretrained, **kwargs)
 
 
 @register_model
-def regnetx_008(pretrained=False, **kwargs):
+def regnetx_008(pretrained=False, **kwargs) -> RegNet:
     """RegNetX-800MF"""
     return _create_regnet('regnetx_008', pretrained, **kwargs)
 
 
 @register_model
-def regnetx_016(pretrained=False, **kwargs):
+def regnetx_016(pretrained=False, **kwargs) -> RegNet:
     """RegNetX-1.6GF"""
     return _create_regnet('regnetx_016', pretrained, **kwargs)
 
 
 @register_model
-def regnetx_032(pretrained=False, **kwargs):
+def regnetx_032(pretrained=False, **kwargs) -> RegNet:
     """RegNetX-3.2GF"""
     return _create_regnet('regnetx_032', pretrained, **kwargs)
 
 
 @register_model
-def regnetx_040(pretrained=False, **kwargs):
+def regnetx_040(pretrained=False, **kwargs) -> RegNet:
     """RegNetX-4.0GF"""
     return _create_regnet('regnetx_040', pretrained, **kwargs)
 
 
 @register_model
-def regnetx_064(pretrained=False, **kwargs):
+def regnetx_064(pretrained=False, **kwargs) -> RegNet:
     """RegNetX-6.4GF"""
     return _create_regnet('regnetx_064', pretrained, **kwargs)
 
 
 @register_model
-def regnetx_080(pretrained=False, **kwargs):
+def regnetx_080(pretrained=False, **kwargs) -> RegNet:
     """RegNetX-8.0GF"""
     return _create_regnet('regnetx_080', pretrained, **kwargs)
 
 
 @register_model
-def regnetx_120(pretrained=False, **kwargs):
+def regnetx_120(pretrained=False, **kwargs) -> RegNet:
     """RegNetX-12GF"""
     return _create_regnet('regnetx_120', pretrained, **kwargs)
 
 
 @register_model
-def regnetx_160(pretrained=False, **kwargs):
+def regnetx_160(pretrained=False, **kwargs) -> RegNet:
     """RegNetX-16GF"""
     return _create_regnet('regnetx_160', pretrained, **kwargs)
 
 
 @register_model
-def regnetx_320(pretrained=False, **kwargs):
+def regnetx_320(pretrained=False, **kwargs) -> RegNet:
     """RegNetX-32GF"""
     return _create_regnet('regnetx_320', pretrained, **kwargs)
 
 
 @register_model
-def regnety_002(pretrained=False, **kwargs):
+def regnety_002(pretrained=False, **kwargs) -> RegNet:
     """RegNetY-200MF"""
     return _create_regnet('regnety_002', pretrained, **kwargs)
 
 
 @register_model
-def regnety_004(pretrained=False, **kwargs):
+def regnety_004(pretrained=False, **kwargs) -> RegNet:
     """RegNetY-400MF"""
     return _create_regnet('regnety_004', pretrained, **kwargs)
 
 
 @register_model
-def regnety_006(pretrained=False, **kwargs):
+def regnety_006(pretrained=False, **kwargs) -> RegNet:
     """RegNetY-600MF"""
     return _create_regnet('regnety_006', pretrained, **kwargs)
 
 
 @register_model
-def regnety_008(pretrained=False, **kwargs):
+def regnety_008(pretrained=False, **kwargs) -> RegNet:
     """RegNetY-800MF"""
     return _create_regnet('regnety_008', pretrained, **kwargs)
 
 
 @register_model
-def regnety_008_tv(pretrained=False, **kwargs):
+def regnety_008_tv(pretrained=False, **kwargs) -> RegNet:
     """RegNetY-800MF w/ torchvision group rounding"""
     return _create_regnet('regnety_008_tv', pretrained, **kwargs)
 
 
 @register_model
-def regnety_016(pretrained=False, **kwargs):
+def regnety_016(pretrained=False, **kwargs) -> RegNet:
     """RegNetY-1.6GF"""
     return _create_regnet('regnety_016', pretrained, **kwargs)
 
 
 @register_model
-def regnety_032(pretrained=False, **kwargs):
+def regnety_032(pretrained=False, **kwargs) -> RegNet:
     """RegNetY-3.2GF"""
     return _create_regnet('regnety_032', pretrained, **kwargs)
 
 
 @register_model
-def regnety_040(pretrained=False, **kwargs):
+def regnety_040(pretrained=False, **kwargs) -> RegNet:
     """RegNetY-4.0GF"""
     return _create_regnet('regnety_040', pretrained, **kwargs)
 
 
 @register_model
-def regnety_064(pretrained=False, **kwargs):
+def regnety_064(pretrained=False, **kwargs) -> RegNet:
     """RegNetY-6.4GF"""
     return _create_regnet('regnety_064', pretrained, **kwargs)
 
 
 @register_model
-def regnety_080(pretrained=False, **kwargs):
+def regnety_080(pretrained=False, **kwargs) -> RegNet:
     """RegNetY-8.0GF"""
     return _create_regnet('regnety_080', pretrained, **kwargs)
 
 
 @register_model
-def regnety_080_tv(pretrained=False, **kwargs):
+def regnety_080_tv(pretrained=False, **kwargs) -> RegNet:
     """RegNetY-8.0GF w/ torchvision group rounding"""
     return _create_regnet('regnety_080_tv', pretrained, **kwargs)
 
 
 @register_model
-def regnety_120(pretrained=False, **kwargs):
+def regnety_120(pretrained=False, **kwargs) -> RegNet:
     """RegNetY-12GF"""
     return _create_regnet('regnety_120', pretrained, **kwargs)
 
 
 @register_model
-def regnety_160(pretrained=False, **kwargs):
+def regnety_160(pretrained=False, **kwargs) -> RegNet:
     """RegNetY-16GF"""
     return _create_regnet('regnety_160', pretrained, **kwargs)
 
 
 @register_model
-def regnety_320(pretrained=False, **kwargs):
+def regnety_320(pretrained=False, **kwargs) -> RegNet:
     """RegNetY-32GF"""
     return _create_regnet('regnety_320', pretrained, **kwargs)
 
 
 @register_model
-def regnety_640(pretrained=False, **kwargs):
+def regnety_640(pretrained=False, **kwargs) -> RegNet:
     """RegNetY-64GF"""
     return _create_regnet('regnety_640', pretrained, **kwargs)
 
 
 @register_model
-def regnety_1280(pretrained=False, **kwargs):
+def regnety_1280(pretrained=False, **kwargs) -> RegNet:
     """RegNetY-128GF"""
     return _create_regnet('regnety_1280', pretrained, **kwargs)
 
 
 @register_model
-def regnety_2560(pretrained=False, **kwargs):
+def regnety_2560(pretrained=False, **kwargs) -> RegNet:
     """RegNetY-256GF"""
     return _create_regnet('regnety_2560', pretrained, **kwargs)
 
 
 @register_model
-def regnety_040_sgn(pretrained=False, **kwargs):
+def regnety_040_sgn(pretrained=False, **kwargs) -> RegNet:
     """RegNetY-4.0GF w/ GroupNorm """
     return _create_regnet('regnety_040_sgn', pretrained, **kwargs)
 
 
 @register_model
-def regnetv_040(pretrained=False, **kwargs):
+def regnetv_040(pretrained=False, **kwargs) -> RegNet:
     """RegNetV-4.0GF (pre-activation)"""
     return _create_regnet('regnetv_040', pretrained, **kwargs)
 
 
 @register_model
-def regnetv_064(pretrained=False, **kwargs):
+def regnetv_064(pretrained=False, **kwargs) -> RegNet:
     """RegNetV-6.4GF (pre-activation)"""
     return _create_regnet('regnetv_064', pretrained, **kwargs)
 
 
 @register_model
-def regnetz_005(pretrained=False, **kwargs):
+def regnetz_005(pretrained=False, **kwargs) -> RegNet:
     """RegNetZ-500MF
     NOTE: config found in https://github.com/facebookresearch/ClassyVision/blob/main/classy_vision/models/regnet.py
     but it's not clear it is equivalent to paper model as not detailed in the paper.
@@ -1100,7 +1100,7 @@ def regnetz_005(pretrained=False, **kwargs):
 
 
 @register_model
-def regnetz_040(pretrained=False, **kwargs):
+def regnetz_040(pretrained=False, **kwargs) -> RegNet:
     """RegNetZ-4.0GF
     NOTE: config found in https://github.com/facebookresearch/ClassyVision/blob/main/classy_vision/models/regnet.py
     but it's not clear it is equivalent to paper model as not detailed in the paper.
@@ -1109,7 +1109,7 @@ def regnetz_040(pretrained=False, **kwargs):
 
 
 @register_model
-def regnetz_040_h(pretrained=False, **kwargs):
+def regnetz_040_h(pretrained=False, **kwargs) -> RegNet:
     """RegNetZ-4.0GF
     NOTE: config found in https://github.com/facebookresearch/ClassyVision/blob/main/classy_vision/models/regnet.py
     but it's not clear it is equivalent to paper model as not detailed in the paper.
diff --git a/timm/models/res2net.py b/timm/models/res2net.py
index 5804a4e8b0..691f929b91 100644
--- a/timm/models/res2net.py
+++ b/timm/models/res2net.py
@@ -145,7 +145,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def res2net50_26w_4s(pretrained=False, **kwargs):
+def res2net50_26w_4s(pretrained=False, **kwargs) -> ResNet:
     """Constructs a Res2Net-50 26w4s model.
     """
     model_args = dict(
@@ -154,7 +154,7 @@ def res2net50_26w_4s(pretrained=False, **kwargs):
 
 
 @register_model
-def res2net101_26w_4s(pretrained=False, **kwargs):
+def res2net101_26w_4s(pretrained=False, **kwargs) -> ResNet:
     """Constructs a Res2Net-101 26w4s model.
     """
     model_args = dict(
@@ -163,7 +163,7 @@ def res2net101_26w_4s(pretrained=False, **kwargs):
 
 
 @register_model
-def res2net50_26w_6s(pretrained=False, **kwargs):
+def res2net50_26w_6s(pretrained=False, **kwargs) -> ResNet:
     """Constructs a Res2Net-50 26w6s model.
     """
     model_args = dict(
@@ -172,7 +172,7 @@ def res2net50_26w_6s(pretrained=False, **kwargs):
 
 
 @register_model
-def res2net50_26w_8s(pretrained=False, **kwargs):
+def res2net50_26w_8s(pretrained=False, **kwargs) -> ResNet:
     """Constructs a Res2Net-50 26w8s model.
     """
     model_args = dict(
@@ -181,7 +181,7 @@ def res2net50_26w_8s(pretrained=False, **kwargs):
 
 
 @register_model
-def res2net50_48w_2s(pretrained=False, **kwargs):
+def res2net50_48w_2s(pretrained=False, **kwargs) -> ResNet:
     """Constructs a Res2Net-50 48w2s model.
     """
     model_args = dict(
@@ -190,7 +190,7 @@ def res2net50_48w_2s(pretrained=False, **kwargs):
 
 
 @register_model
-def res2net50_14w_8s(pretrained=False, **kwargs):
+def res2net50_14w_8s(pretrained=False, **kwargs) -> ResNet:
     """Constructs a Res2Net-50 14w8s model.
     """
     model_args = dict(
@@ -199,7 +199,7 @@ def res2net50_14w_8s(pretrained=False, **kwargs):
 
 
 @register_model
-def res2next50(pretrained=False, **kwargs):
+def res2next50(pretrained=False, **kwargs) -> ResNet:
     """Construct Res2NeXt-50 4s
     """
     model_args = dict(
@@ -208,7 +208,7 @@ def res2next50(pretrained=False, **kwargs):
 
 
 @register_model
-def res2net50d(pretrained=False, **kwargs):
+def res2net50d(pretrained=False, **kwargs) -> ResNet:
     """Construct Res2Net-50
     """
     model_args = dict(
@@ -218,7 +218,7 @@ def res2net50d(pretrained=False, **kwargs):
 
 
 @register_model
-def res2net101d(pretrained=False, **kwargs):
+def res2net101d(pretrained=False, **kwargs) -> ResNet:
     """Construct Res2Net-50
     """
     model_args = dict(
diff --git a/timm/models/resnest.py b/timm/models/resnest.py
index add9afde08..5b1438017e 100644
--- a/timm/models/resnest.py
+++ b/timm/models/resnest.py
@@ -160,7 +160,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def resnest14d(pretrained=False, **kwargs):
+def resnest14d(pretrained=False, **kwargs) -> ResNet:
     """ ResNeSt-14d model. Weights ported from GluonCV.
     """
     model_kwargs = dict(
@@ -171,7 +171,7 @@ def resnest14d(pretrained=False, **kwargs):
 
 
 @register_model
-def resnest26d(pretrained=False, **kwargs):
+def resnest26d(pretrained=False, **kwargs) -> ResNet:
     """ ResNeSt-26d model. Weights ported from GluonCV.
     """
     model_kwargs = dict(
@@ -182,7 +182,7 @@ def resnest26d(pretrained=False, **kwargs):
 
 
 @register_model
-def resnest50d(pretrained=False, **kwargs):
+def resnest50d(pretrained=False, **kwargs) -> ResNet:
     """ ResNeSt-50d model. Matches paper ResNeSt-50 model, https://arxiv.org/abs/2004.08955
     Since this codebase supports all possible variations, 'd' for deep stem, stem_width 32, avg in downsample.
     """
@@ -194,7 +194,7 @@ def resnest50d(pretrained=False, **kwargs):
 
 
 @register_model
-def resnest101e(pretrained=False, **kwargs):
+def resnest101e(pretrained=False, **kwargs) -> ResNet:
     """ ResNeSt-101e model. Matches paper ResNeSt-101 model, https://arxiv.org/abs/2004.08955
      Since this codebase supports all possible variations, 'e' for deep stem, stem_width 64, avg in downsample.
     """
@@ -206,7 +206,7 @@ def resnest101e(pretrained=False, **kwargs):
 
 
 @register_model
-def resnest200e(pretrained=False, **kwargs):
+def resnest200e(pretrained=False, **kwargs) -> ResNet:
     """ ResNeSt-200e model. Matches paper ResNeSt-200 model, https://arxiv.org/abs/2004.08955
     Since this codebase supports all possible variations, 'e' for deep stem, stem_width 64, avg in downsample.
     """
@@ -218,7 +218,7 @@ def resnest200e(pretrained=False, **kwargs):
 
 
 @register_model
-def resnest269e(pretrained=False, **kwargs):
+def resnest269e(pretrained=False, **kwargs) -> ResNet:
     """ ResNeSt-269e model. Matches paper ResNeSt-269 model, https://arxiv.org/abs/2004.08955
     Since this codebase supports all possible variations, 'e' for deep stem, stem_width 64, avg in downsample.
     """
@@ -230,7 +230,7 @@ def resnest269e(pretrained=False, **kwargs):
 
 
 @register_model
-def resnest50d_4s2x40d(pretrained=False, **kwargs):
+def resnest50d_4s2x40d(pretrained=False, **kwargs) -> ResNet:
     """ResNeSt-50 4s2x40d from https://github.com/zhanghang1989/ResNeSt/blob/master/ablation.md
     """
     model_kwargs = dict(
@@ -241,7 +241,7 @@ def resnest50d_4s2x40d(pretrained=False, **kwargs):
 
 
 @register_model
-def resnest50d_1s4x24d(pretrained=False, **kwargs):
+def resnest50d_1s4x24d(pretrained=False, **kwargs) -> ResNet:
     """ResNeSt-50 1s4x24d from https://github.com/zhanghang1989/ResNeSt/blob/master/ablation.md
     """
     model_kwargs = dict(
diff --git a/timm/models/resnet.py b/timm/models/resnet.py
index 18df052a7f..b6e29757ee 100644
--- a/timm/models/resnet.py
+++ b/timm/models/resnet.py
@@ -19,7 +19,7 @@
     get_act_layer, get_norm_layer, create_classifier
 from ._builder import build_model_with_cfg
 from ._manipulate import checkpoint_seq
-from ._registry import register_model, model_entrypoint, generate_default_cfgs, register_model_deprecations
+from ._registry import register_model, generate_default_cfgs, register_model_deprecations
 
 __all__ = ['ResNet', 'BasicBlock', 'Bottleneck']  # model_registry will add each entrypoint fn to this
 
@@ -489,11 +489,6 @@ def __init__(
 
         self.init_weights(zero_init_last=zero_init_last)
 
-    @staticmethod
-    def from_pretrained(model_name: str, load_weights=True, **kwargs) -> 'ResNet':
-        entry_fn = model_entrypoint(model_name, 'resnet')
-        return entry_fn(pretrained=not load_weights, **kwargs)
-
     @torch.jit.ignore
     def init_weights(self, zero_init_last=True):
         for n, m in self.named_modules():
@@ -1202,7 +1197,7 @@ def _gcfg(url='', **kwargs):
 
 
 @register_model
-def resnet10t(pretrained=False, **kwargs):
+def resnet10t(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-10-T model.
     """
     model_args = dict(block=BasicBlock, layers=[1, 1, 1, 1], stem_width=32, stem_type='deep_tiered', avg_down=True)
@@ -1210,7 +1205,7 @@ def resnet10t(pretrained=False, **kwargs):
 
 
 @register_model
-def resnet14t(pretrained=False, **kwargs):
+def resnet14t(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-14-T model.
     """
     model_args = dict(block=Bottleneck, layers=[1, 1, 1, 1], stem_width=32, stem_type='deep_tiered', avg_down=True)
@@ -1218,7 +1213,7 @@ def resnet14t(pretrained=False, **kwargs):
 
 
 @register_model
-def resnet18(pretrained=False, **kwargs):
+def resnet18(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-18 model.
     """
     model_args = dict(block=BasicBlock, layers=[2, 2, 2, 2])
@@ -1226,7 +1221,7 @@ def resnet18(pretrained=False, **kwargs):
 
 
 @register_model
-def resnet18d(pretrained=False, **kwargs):
+def resnet18d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-18-D model.
     """
     model_args = dict(block=BasicBlock, layers=[2, 2, 2, 2], stem_width=32, stem_type='deep', avg_down=True)
@@ -1234,7 +1229,7 @@ def resnet18d(pretrained=False, **kwargs):
 
 
 @register_model
-def resnet34(pretrained=False, **kwargs):
+def resnet34(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-34 model.
     """
     model_args = dict(block=BasicBlock, layers=[3, 4, 6, 3])
@@ -1242,7 +1237,7 @@ def resnet34(pretrained=False, **kwargs):
 
 
 @register_model
-def resnet34d(pretrained=False, **kwargs):
+def resnet34d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-34-D model.
     """
     model_args = dict(block=BasicBlock, layers=[3, 4, 6, 3], stem_width=32, stem_type='deep', avg_down=True)
@@ -1250,7 +1245,7 @@ def resnet34d(pretrained=False, **kwargs):
 
 
 @register_model
-def resnet26(pretrained=False, **kwargs):
+def resnet26(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-26 model.
     """
     model_args = dict(block=Bottleneck, layers=[2, 2, 2, 2])
@@ -1258,7 +1253,7 @@ def resnet26(pretrained=False, **kwargs):
 
 
 @register_model
-def resnet26t(pretrained=False, **kwargs):
+def resnet26t(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-26-T model.
     """
     model_args = dict(block=Bottleneck, layers=[2, 2, 2, 2], stem_width=32, stem_type='deep_tiered', avg_down=True)
@@ -1266,7 +1261,7 @@ def resnet26t(pretrained=False, **kwargs):
 
 
 @register_model
-def resnet26d(pretrained=False, **kwargs):
+def resnet26d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-26-D model.
     """
     model_args = dict(block=Bottleneck, layers=[2, 2, 2, 2], stem_width=32, stem_type='deep', avg_down=True)
@@ -1274,7 +1269,7 @@ def resnet26d(pretrained=False, **kwargs):
 
 
 @register_model
-def resnet50(pretrained=False, **kwargs):
+def resnet50(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-50 model.
     """
     model_args = dict(block=Bottleneck, layers=[3, 4, 6, 3],  **kwargs)
@@ -1306,7 +1301,7 @@ def resnet50s(pretrained=False, **kwargs) -> ResNet:
 
 
 @register_model
-def resnet50t(pretrained=False, **kwargs):
+def resnet50t(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-50-T model.
     """
     model_args = dict(block=Bottleneck, layers=[3, 4, 6, 3], stem_width=32, stem_type='deep_tiered', avg_down=True)
@@ -1314,7 +1309,7 @@ def resnet50t(pretrained=False, **kwargs):
 
 
 @register_model
-def resnet101(pretrained=False, **kwargs):
+def resnet101(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-101 model.
     """
     model_args = dict(block=Bottleneck, layers=[3, 4, 23, 3])
@@ -1322,7 +1317,7 @@ def resnet101(pretrained=False, **kwargs):
 
 
 @register_model
-def resnet101c(pretrained=False, **kwargs):
+def resnet101c(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-101-C model.
     """
     model_args = dict(block=Bottleneck, layers=[3, 4, 23, 3], stem_width=32, stem_type='deep')
@@ -1330,7 +1325,7 @@ def resnet101c(pretrained=False, **kwargs):
 
 
 @register_model
-def resnet101d(pretrained=False, **kwargs):
+def resnet101d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-101-D model.
     """
     model_args = dict(block=Bottleneck, layers=[3, 4, 23, 3], stem_width=32, stem_type='deep', avg_down=True)
@@ -1338,7 +1333,7 @@ def resnet101d(pretrained=False, **kwargs):
 
 
 @register_model
-def resnet101s(pretrained=False, **kwargs):
+def resnet101s(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-101-S model.
     """
     model_args = dict(block=Bottleneck, layers=[3, 4, 23, 3], stem_width=64, stem_type='deep')
@@ -1346,7 +1341,7 @@ def resnet101s(pretrained=False, **kwargs):
 
 
 @register_model
-def resnet152(pretrained=False, **kwargs):
+def resnet152(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-152 model.
     """
     model_args = dict(block=Bottleneck, layers=[3, 8, 36, 3])
@@ -1354,7 +1349,7 @@ def resnet152(pretrained=False, **kwargs):
 
 
 @register_model
-def resnet152c(pretrained=False, **kwargs):
+def resnet152c(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-152-C model.
     """
     model_args = dict(block=Bottleneck, layers=[3, 8, 36, 3], stem_width=32, stem_type='deep')
@@ -1362,7 +1357,7 @@ def resnet152c(pretrained=False, **kwargs):
 
 
 @register_model
-def resnet152d(pretrained=False, **kwargs):
+def resnet152d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-152-D model.
     """
     model_args = dict(block=Bottleneck, layers=[3, 8, 36, 3], stem_width=32, stem_type='deep', avg_down=True)
@@ -1370,7 +1365,7 @@ def resnet152d(pretrained=False, **kwargs):
 
 
 @register_model
-def resnet152s(pretrained=False, **kwargs):
+def resnet152s(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-152-S model.
     """
     model_args = dict(block=Bottleneck, layers=[3, 8, 36, 3], stem_width=64, stem_type='deep')
@@ -1378,7 +1373,7 @@ def resnet152s(pretrained=False, **kwargs):
 
 
 @register_model
-def resnet200(pretrained=False, **kwargs):
+def resnet200(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-200 model.
     """
     model_args = dict(block=Bottleneck, layers=[3, 24, 36, 3])
@@ -1386,7 +1381,7 @@ def resnet200(pretrained=False, **kwargs):
 
 
 @register_model
-def resnet200d(pretrained=False, **kwargs):
+def resnet200d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-200-D model.
     """
     model_args = dict(block=Bottleneck, layers=[3, 24, 36, 3], stem_width=32, stem_type='deep', avg_down=True)
@@ -1394,7 +1389,7 @@ def resnet200d(pretrained=False, **kwargs):
 
 
 @register_model
-def wide_resnet50_2(pretrained=False, **kwargs):
+def wide_resnet50_2(pretrained=False, **kwargs) -> ResNet:
     """Constructs a Wide ResNet-50-2 model.
     The model is the same as ResNet except for the bottleneck number of channels
     which is twice larger in every block. The number of channels in outer 1x1
@@ -1406,7 +1401,7 @@ def wide_resnet50_2(pretrained=False, **kwargs):
 
 
 @register_model
-def wide_resnet101_2(pretrained=False, **kwargs):
+def wide_resnet101_2(pretrained=False, **kwargs) -> ResNet:
     """Constructs a Wide ResNet-101-2 model.
     The model is the same as ResNet except for the bottleneck number of channels
     which is twice larger in every block. The number of channels in outer 1x1
@@ -1417,7 +1412,7 @@ def wide_resnet101_2(pretrained=False, **kwargs):
 
 
 @register_model
-def resnet50_gn(pretrained=False, **kwargs):
+def resnet50_gn(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-50 model w/ GroupNorm
     """
     model_args = dict(block=Bottleneck, layers=[3, 4, 6, 3],  **kwargs)
@@ -1425,7 +1420,7 @@ def resnet50_gn(pretrained=False, **kwargs):
 
 
 @register_model
-def resnext50_32x4d(pretrained=False, **kwargs):
+def resnext50_32x4d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNeXt50-32x4d model.
     """
     model_args = dict(block=Bottleneck, layers=[3, 4, 6, 3], cardinality=32, base_width=4)
@@ -1433,7 +1428,7 @@ def resnext50_32x4d(pretrained=False, **kwargs):
 
 
 @register_model
-def resnext50d_32x4d(pretrained=False, **kwargs):
+def resnext50d_32x4d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNeXt50d-32x4d model. ResNext50 w/ deep stem & avg pool downsample
     """
     model_args = dict(
@@ -1443,7 +1438,7 @@ def resnext50d_32x4d(pretrained=False, **kwargs):
 
 
 @register_model
-def resnext101_32x4d(pretrained=False, **kwargs):
+def resnext101_32x4d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNeXt-101 32x4d model.
     """
     model_args = dict(block=Bottleneck, layers=[3, 4, 23, 3], cardinality=32, base_width=4)
@@ -1451,7 +1446,7 @@ def resnext101_32x4d(pretrained=False, **kwargs):
 
 
 @register_model
-def resnext101_32x8d(pretrained=False, **kwargs):
+def resnext101_32x8d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNeXt-101 32x8d model.
     """
     model_args = dict(block=Bottleneck, layers=[3, 4, 23, 3], cardinality=32, base_width=8)
@@ -1459,7 +1454,7 @@ def resnext101_32x8d(pretrained=False, **kwargs):
 
 
 @register_model
-def resnext101_32x16d(pretrained=False, **kwargs):
+def resnext101_32x16d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNeXt-101 32x16d model
     """
     model_args = dict(block=Bottleneck, layers=[3, 4, 23, 3], cardinality=32, base_width=16)
@@ -1467,7 +1462,7 @@ def resnext101_32x16d(pretrained=False, **kwargs):
 
 
 @register_model
-def resnext101_32x32d(pretrained=False, **kwargs):
+def resnext101_32x32d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNeXt-101 32x32d model
     """
     model_args = dict(block=Bottleneck, layers=[3, 4, 23, 3], cardinality=32, base_width=32)
@@ -1475,7 +1470,7 @@ def resnext101_32x32d(pretrained=False, **kwargs):
 
 
 @register_model
-def resnext101_64x4d(pretrained=False, **kwargs):
+def resnext101_64x4d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNeXt101-64x4d model.
     """
     model_args = dict(block=Bottleneck, layers=[3, 4, 23, 3], cardinality=64, base_width=4)
@@ -1483,7 +1478,7 @@ def resnext101_64x4d(pretrained=False, **kwargs):
 
 
 @register_model
-def ecaresnet26t(pretrained=False, **kwargs):
+def ecaresnet26t(pretrained=False, **kwargs) -> ResNet:
     """Constructs an ECA-ResNeXt-26-T model.
     This is technically a 28 layer ResNet, like a 'D' bag-of-tricks model but with tiered 24, 32, 64 channels
     in the deep stem and ECA attn.
@@ -1495,7 +1490,7 @@ def ecaresnet26t(pretrained=False, **kwargs):
 
 
 @register_model
-def ecaresnet50d(pretrained=False, **kwargs):
+def ecaresnet50d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-50-D model with eca.
     """
     model_args = dict(
@@ -1505,7 +1500,7 @@ def ecaresnet50d(pretrained=False, **kwargs):
 
 
 @register_model
-def ecaresnet50d_pruned(pretrained=False, **kwargs):
+def ecaresnet50d_pruned(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-50-D model pruned with eca.
         The pruning has been obtained using https://arxiv.org/pdf/2002.08258.pdf
     """
@@ -1516,7 +1511,7 @@ def ecaresnet50d_pruned(pretrained=False, **kwargs):
 
 
 @register_model
-def ecaresnet50t(pretrained=False, **kwargs):
+def ecaresnet50t(pretrained=False, **kwargs) -> ResNet:
     """Constructs an ECA-ResNet-50-T model.
     Like a 'D' bag-of-tricks model but with tiered 24, 32, 64 channels in the deep stem and ECA attn.
     """
@@ -1527,7 +1522,7 @@ def ecaresnet50t(pretrained=False, **kwargs):
 
 
 @register_model
-def ecaresnetlight(pretrained=False, **kwargs):
+def ecaresnetlight(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-50-D light model with eca.
     """
     model_args = dict(
@@ -1537,7 +1532,7 @@ def ecaresnetlight(pretrained=False, **kwargs):
 
 
 @register_model
-def ecaresnet101d(pretrained=False, **kwargs):
+def ecaresnet101d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-101-D model with eca.
     """
     model_args = dict(
@@ -1547,7 +1542,7 @@ def ecaresnet101d(pretrained=False, **kwargs):
 
 
 @register_model
-def ecaresnet101d_pruned(pretrained=False, **kwargs):
+def ecaresnet101d_pruned(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-101-D model pruned with eca.
        The pruning has been obtained using https://arxiv.org/pdf/2002.08258.pdf
     """
@@ -1558,7 +1553,7 @@ def ecaresnet101d_pruned(pretrained=False, **kwargs):
 
 
 @register_model
-def ecaresnet200d(pretrained=False, **kwargs):
+def ecaresnet200d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-200-D model with ECA.
     """
     model_args = dict(
@@ -1568,7 +1563,7 @@ def ecaresnet200d(pretrained=False, **kwargs):
 
 
 @register_model
-def ecaresnet269d(pretrained=False, **kwargs):
+def ecaresnet269d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-269-D model with ECA.
     """
     model_args = dict(
@@ -1578,7 +1573,7 @@ def ecaresnet269d(pretrained=False, **kwargs):
 
 
 @register_model
-def ecaresnext26t_32x4d(pretrained=False, **kwargs):
+def ecaresnext26t_32x4d(pretrained=False, **kwargs) -> ResNet:
     """Constructs an ECA-ResNeXt-26-T model.
     This is technically a 28 layer ResNet, like a 'D' bag-of-tricks model but with tiered 24, 32, 64 channels
     in the deep stem. This model replaces SE module with the ECA module
@@ -1590,7 +1585,7 @@ def ecaresnext26t_32x4d(pretrained=False, **kwargs):
 
 
 @register_model
-def ecaresnext50t_32x4d(pretrained=False, **kwargs):
+def ecaresnext50t_32x4d(pretrained=False, **kwargs) -> ResNet:
     """Constructs an ECA-ResNeXt-50-T model.
     This is technically a 28 layer ResNet, like a 'D' bag-of-tricks model but with tiered 24, 32, 64 channels
     in the deep stem. This model replaces SE module with the ECA module
@@ -1602,25 +1597,25 @@ def ecaresnext50t_32x4d(pretrained=False, **kwargs):
 
 
 @register_model
-def seresnet18(pretrained=False, **kwargs):
+def seresnet18(pretrained=False, **kwargs) -> ResNet:
     model_args = dict(block=BasicBlock, layers=[2, 2, 2, 2], block_args=dict(attn_layer='se'))
     return _create_resnet('seresnet18', pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
-def seresnet34(pretrained=False, **kwargs):
+def seresnet34(pretrained=False, **kwargs) -> ResNet:
     model_args = dict(block=BasicBlock, layers=[3, 4, 6, 3], block_args=dict(attn_layer='se'))
     return _create_resnet('seresnet34', pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
-def seresnet50(pretrained=False, **kwargs):
+def seresnet50(pretrained=False, **kwargs) -> ResNet:
     model_args = dict(block=Bottleneck, layers=[3, 4, 6, 3], block_args=dict(attn_layer='se'))
     return _create_resnet('seresnet50', pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
-def seresnet50t(pretrained=False, **kwargs):
+def seresnet50t(pretrained=False, **kwargs) -> ResNet:
     model_args = dict(
         block=Bottleneck, layers=[3, 4, 6, 3],  stem_width=32, stem_type='deep_tiered',
         avg_down=True, block_args=dict(attn_layer='se'))
@@ -1628,19 +1623,19 @@ def seresnet50t(pretrained=False, **kwargs):
 
 
 @register_model
-def seresnet101(pretrained=False, **kwargs):
+def seresnet101(pretrained=False, **kwargs) -> ResNet:
     model_args = dict(block=Bottleneck, layers=[3, 4, 23, 3], block_args=dict(attn_layer='se'))
     return _create_resnet('seresnet101', pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
-def seresnet152(pretrained=False, **kwargs):
+def seresnet152(pretrained=False, **kwargs) -> ResNet:
     model_args = dict(block=Bottleneck, layers=[3, 8, 36, 3], block_args=dict(attn_layer='se'))
     return _create_resnet('seresnet152', pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
-def seresnet152d(pretrained=False, **kwargs):
+def seresnet152d(pretrained=False, **kwargs) -> ResNet:
     model_args = dict(
         block=Bottleneck, layers=[3, 8, 36, 3], stem_width=32, stem_type='deep',
         avg_down=True, block_args=dict(attn_layer='se'))
@@ -1648,7 +1643,7 @@ def seresnet152d(pretrained=False, **kwargs):
 
 
 @register_model
-def seresnet200d(pretrained=False, **kwargs):
+def seresnet200d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-200-D model with SE attn.
     """
     model_args = dict(
@@ -1658,7 +1653,7 @@ def seresnet200d(pretrained=False, **kwargs):
 
 
 @register_model
-def seresnet269d(pretrained=False, **kwargs):
+def seresnet269d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-269-D model with SE attn.
     """
     model_args = dict(
@@ -1668,7 +1663,7 @@ def seresnet269d(pretrained=False, **kwargs):
 
 
 @register_model
-def seresnext26d_32x4d(pretrained=False, **kwargs):
+def seresnext26d_32x4d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a SE-ResNeXt-26-D model.`
     This is technically a 28 layer ResNet, using the 'D' modifier from Gluon / bag-of-tricks for
     combination of deep stem and avg_pool in downsample.
@@ -1680,7 +1675,7 @@ def seresnext26d_32x4d(pretrained=False, **kwargs):
 
 
 @register_model
-def seresnext26t_32x4d(pretrained=False, **kwargs):
+def seresnext26t_32x4d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a SE-ResNet-26-T model.
     This is technically a 28 layer ResNet, like a 'D' bag-of-tricks model but with tiered 24, 32, 64 channels
     in the deep stem.
@@ -1692,7 +1687,7 @@ def seresnext26t_32x4d(pretrained=False, **kwargs):
 
 
 @register_model
-def seresnext26tn_32x4d(pretrained=False, **kwargs):
+def seresnext26tn_32x4d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a SE-ResNeXt-26-T model.
     NOTE I deprecated previous 't' model defs and replaced 't' with 'tn', this was the only tn model of note
     so keeping this def for backwards compat with any uses out there. Old 't' model is lost.
@@ -1701,7 +1696,7 @@ def seresnext26tn_32x4d(pretrained=False, **kwargs):
 
 
 @register_model
-def seresnext50_32x4d(pretrained=False, **kwargs):
+def seresnext50_32x4d(pretrained=False, **kwargs) -> ResNet:
     model_args = dict(
         block=Bottleneck, layers=[3, 4, 6, 3], cardinality=32, base_width=4,
         block_args=dict(attn_layer='se'))
@@ -1709,7 +1704,7 @@ def seresnext50_32x4d(pretrained=False, **kwargs):
 
 
 @register_model
-def seresnext101_32x4d(pretrained=False, **kwargs):
+def seresnext101_32x4d(pretrained=False, **kwargs) -> ResNet:
     model_args = dict(
         block=Bottleneck, layers=[3, 4, 23, 3], cardinality=32, base_width=4,
         block_args=dict(attn_layer='se'))
@@ -1717,7 +1712,7 @@ def seresnext101_32x4d(pretrained=False, **kwargs):
 
 
 @register_model
-def seresnext101_32x8d(pretrained=False, **kwargs):
+def seresnext101_32x8d(pretrained=False, **kwargs) -> ResNet:
     model_args = dict(
         block=Bottleneck, layers=[3, 4, 23, 3], cardinality=32, base_width=8,
         block_args=dict(attn_layer='se'))
@@ -1725,7 +1720,7 @@ def seresnext101_32x8d(pretrained=False, **kwargs):
 
 
 @register_model
-def seresnext101d_32x8d(pretrained=False, **kwargs):
+def seresnext101d_32x8d(pretrained=False, **kwargs) -> ResNet:
     model_args = dict(
         block=Bottleneck, layers=[3, 4, 23, 3], cardinality=32, base_width=8,
         stem_width=32, stem_type='deep', avg_down=True,
@@ -1734,7 +1729,7 @@ def seresnext101d_32x8d(pretrained=False, **kwargs):
 
 
 @register_model
-def seresnext101_64x4d(pretrained=False, **kwargs):
+def seresnext101_64x4d(pretrained=False, **kwargs) -> ResNet:
     model_args = dict(
         block=Bottleneck, layers=[3, 4, 23, 3], cardinality=64, base_width=4,
         block_args=dict(attn_layer='se'))
@@ -1742,7 +1737,7 @@ def seresnext101_64x4d(pretrained=False, **kwargs):
 
 
 @register_model
-def senet154(pretrained=False, **kwargs):
+def senet154(pretrained=False, **kwargs) -> ResNet:
     model_args = dict(
         block=Bottleneck, layers=[3, 8, 36, 3], cardinality=64, base_width=4, stem_type='deep',
         down_kernel_size=3, block_reduce_first=2, block_args=dict(attn_layer='se'))
@@ -1750,7 +1745,7 @@ def senet154(pretrained=False, **kwargs):
 
 
 @register_model
-def resnetblur18(pretrained=False, **kwargs):
+def resnetblur18(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-18 model with blur anti-aliasing
     """
     model_args = dict(block=BasicBlock, layers=[2, 2, 2, 2], aa_layer=BlurPool2d)
@@ -1758,7 +1753,7 @@ def resnetblur18(pretrained=False, **kwargs):
 
 
 @register_model
-def resnetblur50(pretrained=False, **kwargs):
+def resnetblur50(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-50 model with blur anti-aliasing
     """
     model_args = dict(block=Bottleneck, layers=[3, 4, 6, 3], aa_layer=BlurPool2d)
@@ -1766,7 +1761,7 @@ def resnetblur50(pretrained=False, **kwargs):
 
 
 @register_model
-def resnetblur50d(pretrained=False, **kwargs):
+def resnetblur50d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-50-D model with blur anti-aliasing
     """
     model_args = dict(
@@ -1776,7 +1771,7 @@ def resnetblur50d(pretrained=False, **kwargs):
 
 
 @register_model
-def resnetblur101d(pretrained=False, **kwargs):
+def resnetblur101d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-101-D model with blur anti-aliasing
     """
     model_args = dict(
@@ -1786,7 +1781,7 @@ def resnetblur101d(pretrained=False, **kwargs):
 
 
 @register_model
-def resnetaa34d(pretrained=False, **kwargs):
+def resnetaa34d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-34-D model w/ avgpool anti-aliasing
     """
     model_args = dict(
@@ -1795,7 +1790,7 @@ def resnetaa34d(pretrained=False, **kwargs):
 
 
 @register_model
-def resnetaa50(pretrained=False, **kwargs):
+def resnetaa50(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-50 model with avgpool anti-aliasing
     """
     model_args = dict(block=Bottleneck, layers=[3, 4, 6, 3], aa_layer=nn.AvgPool2d)
@@ -1803,7 +1798,7 @@ def resnetaa50(pretrained=False, **kwargs):
 
 
 @register_model
-def resnetaa50d(pretrained=False, **kwargs):
+def resnetaa50d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-50-D model with avgpool anti-aliasing
     """
     model_args = dict(
@@ -1813,7 +1808,7 @@ def resnetaa50d(pretrained=False, **kwargs):
 
 
 @register_model
-def resnetaa101d(pretrained=False, **kwargs):
+def resnetaa101d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-101-D model with avgpool anti-aliasing
     """
     model_args = dict(
@@ -1823,7 +1818,7 @@ def resnetaa101d(pretrained=False, **kwargs):
 
 
 @register_model
-def seresnetaa50d(pretrained=False, **kwargs):
+def seresnetaa50d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a SE=ResNet-50-D model with avgpool anti-aliasing
     """
     model_args = dict(
@@ -1833,7 +1828,7 @@ def seresnetaa50d(pretrained=False, **kwargs):
 
 
 @register_model
-def seresnextaa101d_32x8d(pretrained=False, **kwargs):
+def seresnextaa101d_32x8d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a SE=ResNeXt-101-D 32x8d model with avgpool anti-aliasing
     """
     model_args = dict(
@@ -1844,7 +1839,7 @@ def seresnextaa101d_32x8d(pretrained=False, **kwargs):
 
 
 @register_model
-def resnetrs50(pretrained=False, **kwargs):
+def resnetrs50(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-RS-50 model.
     Paper: Revisiting ResNets - https://arxiv.org/abs/2103.07579
     Pretrained weights from https://github.com/tensorflow/tpu/tree/bee9c4f6/models/official/resnet/resnet_rs
@@ -1857,7 +1852,7 @@ def resnetrs50(pretrained=False, **kwargs):
 
 
 @register_model
-def resnetrs101(pretrained=False, **kwargs):
+def resnetrs101(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-RS-101 model.
     Paper: Revisiting ResNets - https://arxiv.org/abs/2103.07579
     Pretrained weights from https://github.com/tensorflow/tpu/tree/bee9c4f6/models/official/resnet/resnet_rs
@@ -1870,7 +1865,7 @@ def resnetrs101(pretrained=False, **kwargs):
 
 
 @register_model
-def resnetrs152(pretrained=False, **kwargs):
+def resnetrs152(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-RS-152 model.
     Paper: Revisiting ResNets - https://arxiv.org/abs/2103.07579
     Pretrained weights from https://github.com/tensorflow/tpu/tree/bee9c4f6/models/official/resnet/resnet_rs
@@ -1883,7 +1878,7 @@ def resnetrs152(pretrained=False, **kwargs):
 
 
 @register_model
-def resnetrs200(pretrained=False, **kwargs):
+def resnetrs200(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-RS-200 model.
     Paper: Revisiting ResNets - https://arxiv.org/abs/2103.07579
     Pretrained weights from https://github.com/tensorflow/tpu/tree/bee9c4f6/models/official/resnet/resnet_rs
@@ -1896,7 +1891,7 @@ def resnetrs200(pretrained=False, **kwargs):
 
 
 @register_model
-def resnetrs270(pretrained=False, **kwargs):
+def resnetrs270(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-RS-270 model.
     Paper: Revisiting ResNets - https://arxiv.org/abs/2103.07579
     Pretrained weights from https://github.com/tensorflow/tpu/tree/bee9c4f6/models/official/resnet/resnet_rs
@@ -1910,7 +1905,7 @@ def resnetrs270(pretrained=False, **kwargs):
 
 
 @register_model
-def resnetrs350(pretrained=False, **kwargs):
+def resnetrs350(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-RS-350 model.
     Paper: Revisiting ResNets - https://arxiv.org/abs/2103.07579
     Pretrained weights from https://github.com/tensorflow/tpu/tree/bee9c4f6/models/official/resnet/resnet_rs
@@ -1923,7 +1918,7 @@ def resnetrs350(pretrained=False, **kwargs):
 
 
 @register_model
-def resnetrs420(pretrained=False, **kwargs):
+def resnetrs420(pretrained=False, **kwargs) -> ResNet:
     """Constructs a ResNet-RS-420 model
     Paper: Revisiting ResNets - https://arxiv.org/abs/2103.07579
     Pretrained weights from https://github.com/tensorflow/tpu/tree/bee9c4f6/models/official/resnet/resnet_rs
diff --git a/timm/models/resnetv2.py b/timm/models/resnetv2.py
index 2239889d34..43d2704e42 100644
--- a/timm/models/resnetv2.py
+++ b/timm/models/resnetv2.py
@@ -643,49 +643,49 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def resnetv2_50x1_bit(pretrained=False, **kwargs):
+def resnetv2_50x1_bit(pretrained=False, **kwargs) -> ResNetV2:
     return _create_resnetv2_bit(
         'resnetv2_50x1_bit', pretrained=pretrained, layers=[3, 4, 6, 3], width_factor=1, **kwargs)
 
 
 @register_model
-def resnetv2_50x3_bit(pretrained=False, **kwargs):
+def resnetv2_50x3_bit(pretrained=False, **kwargs) -> ResNetV2:
     return _create_resnetv2_bit(
         'resnetv2_50x3_bit', pretrained=pretrained, layers=[3, 4, 6, 3], width_factor=3, **kwargs)
 
 
 @register_model
-def resnetv2_101x1_bit(pretrained=False, **kwargs):
+def resnetv2_101x1_bit(pretrained=False, **kwargs) -> ResNetV2:
     return _create_resnetv2_bit(
         'resnetv2_101x1_bit', pretrained=pretrained, layers=[3, 4, 23, 3], width_factor=1, **kwargs)
 
 
 @register_model
-def resnetv2_101x3_bit(pretrained=False, **kwargs):
+def resnetv2_101x3_bit(pretrained=False, **kwargs) -> ResNetV2:
     return _create_resnetv2_bit(
         'resnetv2_101x3_bit', pretrained=pretrained, layers=[3, 4, 23, 3], width_factor=3, **kwargs)
 
 
 @register_model
-def resnetv2_152x2_bit(pretrained=False, **kwargs):
+def resnetv2_152x2_bit(pretrained=False, **kwargs) -> ResNetV2:
     return _create_resnetv2_bit(
         'resnetv2_152x2_bit', pretrained=pretrained, layers=[3, 8, 36, 3], width_factor=2, **kwargs)
 
 
 @register_model
-def resnetv2_152x4_bit(pretrained=False, **kwargs):
+def resnetv2_152x4_bit(pretrained=False, **kwargs) -> ResNetV2:
     return _create_resnetv2_bit(
         'resnetv2_152x4_bit', pretrained=pretrained, layers=[3, 8, 36, 3], width_factor=4, **kwargs)
 
 
 @register_model
-def resnetv2_50(pretrained=False, **kwargs):
+def resnetv2_50(pretrained=False, **kwargs) -> ResNetV2:
     model_args = dict(layers=[3, 4, 6, 3], conv_layer=create_conv2d, norm_layer=BatchNormAct2d)
     return _create_resnetv2('resnetv2_50', pretrained=pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
-def resnetv2_50d(pretrained=False, **kwargs):
+def resnetv2_50d(pretrained=False, **kwargs) -> ResNetV2:
     model_args = dict(
         layers=[3, 4, 6, 3], conv_layer=create_conv2d, norm_layer=BatchNormAct2d,
         stem_type='deep', avg_down=True)
@@ -693,7 +693,7 @@ def resnetv2_50d(pretrained=False, **kwargs):
 
 
 @register_model
-def resnetv2_50t(pretrained=False, **kwargs):
+def resnetv2_50t(pretrained=False, **kwargs) -> ResNetV2:
     model_args = dict(
         layers=[3, 4, 6, 3], conv_layer=create_conv2d, norm_layer=BatchNormAct2d,
         stem_type='tiered', avg_down=True)
@@ -701,13 +701,13 @@ def resnetv2_50t(pretrained=False, **kwargs):
 
 
 @register_model
-def resnetv2_101(pretrained=False, **kwargs):
+def resnetv2_101(pretrained=False, **kwargs) -> ResNetV2:
     model_args = dict(layers=[3, 4, 23, 3], conv_layer=create_conv2d, norm_layer=BatchNormAct2d)
     return _create_resnetv2('resnetv2_101', pretrained=pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
-def resnetv2_101d(pretrained=False, **kwargs):
+def resnetv2_101d(pretrained=False, **kwargs) -> ResNetV2:
     model_args = dict(
         layers=[3, 4, 23, 3], conv_layer=create_conv2d, norm_layer=BatchNormAct2d,
         stem_type='deep', avg_down=True)
@@ -715,13 +715,13 @@ def resnetv2_101d(pretrained=False, **kwargs):
 
 
 @register_model
-def resnetv2_152(pretrained=False, **kwargs):
+def resnetv2_152(pretrained=False, **kwargs) -> ResNetV2:
     model_args = dict(layers=[3, 8, 36, 3], conv_layer=create_conv2d, norm_layer=BatchNormAct2d)
     return _create_resnetv2('resnetv2_152', pretrained=pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
-def resnetv2_152d(pretrained=False, **kwargs):
+def resnetv2_152d(pretrained=False, **kwargs) -> ResNetV2:
     model_args = dict(
         layers=[3, 8, 36, 3], conv_layer=create_conv2d, norm_layer=BatchNormAct2d,
         stem_type='deep', avg_down=True)
@@ -731,7 +731,7 @@ def resnetv2_152d(pretrained=False, **kwargs):
 # Experimental configs (may change / be removed)
 
 @register_model
-def resnetv2_50d_gn(pretrained=False, **kwargs):
+def resnetv2_50d_gn(pretrained=False, **kwargs) -> ResNetV2:
     model_args = dict(
         layers=[3, 4, 6, 3], conv_layer=create_conv2d, norm_layer=GroupNormAct,
         stem_type='deep', avg_down=True)
@@ -739,7 +739,7 @@ def resnetv2_50d_gn(pretrained=False, **kwargs):
 
 
 @register_model
-def resnetv2_50d_evos(pretrained=False, **kwargs):
+def resnetv2_50d_evos(pretrained=False, **kwargs) -> ResNetV2:
     model_args = dict(
         layers=[3, 4, 6, 3], conv_layer=create_conv2d, norm_layer=EvoNorm2dS0,
         stem_type='deep', avg_down=True)
@@ -747,7 +747,7 @@ def resnetv2_50d_evos(pretrained=False, **kwargs):
 
 
 @register_model
-def resnetv2_50d_frn(pretrained=False, **kwargs):
+def resnetv2_50d_frn(pretrained=False, **kwargs) -> ResNetV2:
     model_args = dict(
         layers=[3, 4, 6, 3], conv_layer=create_conv2d, norm_layer=FilterResponseNormTlu2d,
         stem_type='deep', avg_down=True)
diff --git a/timm/models/rexnet.py b/timm/models/rexnet.py
index 4fa5c4ba83..d34933f782 100644
--- a/timm/models/rexnet.py
+++ b/timm/models/rexnet.py
@@ -23,7 +23,7 @@
 from ._manipulate import checkpoint_seq
 from ._registry import generate_default_cfgs, register_model
 
-__all__ = ['ReXNetV1']  # model_registry will add each entrypoint fn to this
+__all__ = ['RexNet']  # model_registry will add each entrypoint fn to this
 
 
 SEWithNorm = partial(SEModule, norm_layer=nn.BatchNorm2d)
@@ -167,7 +167,7 @@ def _build_blocks(
     return features, feature_info
 
 
-class ReXNetV1(nn.Module):
+class RexNet(nn.Module):
     def __init__(
             self,
             in_chans=3,
@@ -185,7 +185,7 @@ def __init__(
             drop_rate=0.2,
             drop_path_rate=0.,
     ):
-        super(ReXNetV1, self).__init__()
+        super(RexNet, self).__init__()
         self.num_classes = num_classes
         self.drop_rate = drop_rate
         self.grad_checkpointing = False
@@ -252,7 +252,7 @@ def forward(self, x):
 def _create_rexnet(variant, pretrained, **kwargs):
     feature_cfg = dict(flatten_sequential=True)
     return build_model_with_cfg(
-        ReXNetV1,
+        RexNet,
         variant,
         pretrained,
         feature_cfg=feature_cfg,
@@ -297,60 +297,60 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def rexnet_100(pretrained=False, **kwargs):
+def rexnet_100(pretrained=False, **kwargs) -> RexNet:
     """ReXNet V1 1.0x"""
     return _create_rexnet('rexnet_100', pretrained, **kwargs)
 
 
 @register_model
-def rexnet_130(pretrained=False, **kwargs):
+def rexnet_130(pretrained=False, **kwargs) -> RexNet:
     """ReXNet V1 1.3x"""
     return _create_rexnet('rexnet_130', pretrained, width_mult=1.3, **kwargs)
 
 
 @register_model
-def rexnet_150(pretrained=False, **kwargs):
+def rexnet_150(pretrained=False, **kwargs) -> RexNet:
     """ReXNet V1 1.5x"""
     return _create_rexnet('rexnet_150', pretrained, width_mult=1.5, **kwargs)
 
 
 @register_model
-def rexnet_200(pretrained=False, **kwargs):
+def rexnet_200(pretrained=False, **kwargs) -> RexNet:
     """ReXNet V1 2.0x"""
     return _create_rexnet('rexnet_200', pretrained, width_mult=2.0, **kwargs)
 
 
 @register_model
-def rexnet_300(pretrained=False, **kwargs):
+def rexnet_300(pretrained=False, **kwargs) -> RexNet:
     """ReXNet V1 3.0x"""
     return _create_rexnet('rexnet_300', pretrained, width_mult=3.0, **kwargs)
 
 
 @register_model
-def rexnetr_100(pretrained=False, **kwargs):
+def rexnetr_100(pretrained=False, **kwargs) -> RexNet:
     """ReXNet V1 1.0x w/ rounded (mod 8) channels"""
     return _create_rexnet('rexnetr_100', pretrained, ch_div=8, **kwargs)
 
 
 @register_model
-def rexnetr_130(pretrained=False, **kwargs):
+def rexnetr_130(pretrained=False, **kwargs) -> RexNet:
     """ReXNet V1 1.3x w/ rounded (mod 8) channels"""
     return _create_rexnet('rexnetr_130', pretrained, width_mult=1.3, ch_div=8, **kwargs)
 
 
 @register_model
-def rexnetr_150(pretrained=False, **kwargs):
+def rexnetr_150(pretrained=False, **kwargs) -> RexNet:
     """ReXNet V1 1.5x w/ rounded (mod 8) channels"""
     return _create_rexnet('rexnetr_150', pretrained, width_mult=1.5, ch_div=8, **kwargs)
 
 
 @register_model
-def rexnetr_200(pretrained=False, **kwargs):
+def rexnetr_200(pretrained=False, **kwargs) -> RexNet:
     """ReXNet V1 2.0x w/ rounded (mod 8) channels"""
     return _create_rexnet('rexnetr_200', pretrained, width_mult=2.0, ch_div=8, **kwargs)
 
 
 @register_model
-def rexnetr_300(pretrained=False, **kwargs):
+def rexnetr_300(pretrained=False, **kwargs) -> RexNet:
     """ReXNet V1 3.0x w/ rounded (mod 16) channels"""
     return _create_rexnet('rexnetr_300', pretrained, width_mult=3.0, ch_div=16, **kwargs)
diff --git a/timm/models/selecsls.py b/timm/models/selecsls.py
index 8aeaadd95a..0a333a3d21 100644
--- a/timm/models/selecsls.py
+++ b/timm/models/selecsls.py
@@ -20,7 +20,7 @@
 from ._builder import build_model_with_cfg
 from ._registry import register_model, generate_default_cfgs
 
-__all__ = ['SelecSLS']  # model_registry will add each entrypoint fn to this
+__all__ = ['SelecSls']  # model_registry will add each entrypoint fn to this
 
 
 class SequentialList(nn.Sequential):
@@ -77,9 +77,9 @@ def conv_bn(in_chs, out_chs, k=3, stride=1, padding=None, dilation=1):
     )
 
 
-class SelecSLSBlock(nn.Module):
+class SelecSlsBlock(nn.Module):
     def __init__(self, in_chs, skip_chs, mid_chs, out_chs, is_first, stride, dilation=1):
-        super(SelecSLSBlock, self).__init__()
+        super(SelecSlsBlock, self).__init__()
         self.stride = stride
         self.is_first = is_first
         assert stride in [1, 2]
@@ -107,8 +107,8 @@ def forward(self, x: List[torch.Tensor]) -> List[torch.Tensor]:
             return [self.conv6(torch.cat([d1, d2, d3, x[1]], 1)), x[1]]
 
 
-class SelecSLS(nn.Module):
-    """SelecSLS42 / SelecSLS60 / SelecSLS84
+class SelecSls(nn.Module):
+    """SelecSls42 / SelecSls60 / SelecSls84
 
     Parameters
     ----------
@@ -125,7 +125,7 @@ class SelecSLS(nn.Module):
 
     def __init__(self, cfg, num_classes=1000, in_chans=3, drop_rate=0.0, global_pool='avg'):
         self.num_classes = num_classes
-        super(SelecSLS, self).__init__()
+        super(SelecSls, self).__init__()
 
         self.stem = conv_bn(in_chans, 32, stride=2)
         self.features = SequentialList(*[cfg['block'](*block_args) for block_args in cfg['features']])
@@ -182,11 +182,11 @@ def forward(self, x):
         return x
 
 
-def _create_selecsls(variant, pretrained, **kwargs):
+def _create_SelecSls(variant, pretrained, **kwargs):
     cfg = {}
     feature_info = [dict(num_chs=32, reduction=2, module='stem.2')]
-    if variant.startswith('selecsls42'):
-        cfg['block'] = SelecSLSBlock
+    if variant.startswith('SelecSls42'):
+        cfg['block'] = SelecSlsBlock
         # Define configuration of the network after the initial neck
         cfg['features'] = [
             # in_chs, skip_chs, mid_chs, out_chs, is_first, stride
@@ -204,7 +204,7 @@ def _create_selecsls(variant, pretrained, **kwargs):
         ])
         # Head can be replaced with alternative configurations depending on the problem
         feature_info.append(dict(num_chs=1024, reduction=32, module='head.1'))
-        if variant == 'selecsls42b':
+        if variant == 'SelecSls42b':
             cfg['head'] = [
                 (480, 960, 3, 2),
                 (960, 1024, 3, 1),
@@ -223,8 +223,8 @@ def _create_selecsls(variant, pretrained, **kwargs):
             feature_info.append(dict(num_chs=1280, reduction=64, module='head.3'))
             cfg['num_features'] = 1280
 
-    elif variant.startswith('selecsls60'):
-        cfg['block'] = SelecSLSBlock
+    elif variant.startswith('SelecSls60'):
+        cfg['block'] = SelecSlsBlock
         # Define configuration of the network after the initial neck
         cfg['features'] = [
             # in_chs, skip_chs, mid_chs, out_chs, is_first, stride
@@ -245,7 +245,7 @@ def _create_selecsls(variant, pretrained, **kwargs):
         ])
         # Head can be replaced with alternative configurations depending on the problem
         feature_info.append(dict(num_chs=1024, reduction=32, module='head.1'))
-        if variant == 'selecsls60b':
+        if variant == 'SelecSls60b':
             cfg['head'] = [
                 (416, 756, 3, 2),
                 (756, 1024, 3, 1),
@@ -264,8 +264,8 @@ def _create_selecsls(variant, pretrained, **kwargs):
             feature_info.append(dict(num_chs=1280, reduction=64, module='head.3'))
             cfg['num_features'] = 1280
 
-    elif variant == 'selecsls84':
-        cfg['block'] = SelecSLSBlock
+    elif variant == 'SelecSls84':
+        cfg['block'] = SelecSlsBlock
         # Define configuration of the network after the initial neck
         cfg['features'] = [
             # in_chs, skip_chs, mid_chs, out_chs, is_first, stride
@@ -306,7 +306,7 @@ def _create_selecsls(variant, pretrained, **kwargs):
 
     # this model can do 6 feature levels by default, unlike most others, leave as 0-4 to avoid surprises?
     return build_model_with_cfg(
-        SelecSLS,
+        SelecSls,
         variant,
         pretrained,
         model_cfg=cfg,
@@ -327,52 +327,52 @@ def _cfg(url='', **kwargs):
 
 
 default_cfgs = generate_default_cfgs({
-    'selecsls42.untrained': _cfg(
+    'SelecSls42.untrained': _cfg(
         interpolation='bicubic'),
-    'selecsls42b.in1k': _cfg(
+    'SelecSls42b.in1k': _cfg(
         hf_hub_id='timm/',
         interpolation='bicubic'),
-    'selecsls60.in1k': _cfg(
+    'SelecSls60.in1k': _cfg(
         hf_hub_id='timm/',
         interpolation='bicubic'),
-    'selecsls60b.in1k': _cfg(
+    'SelecSls60b.in1k': _cfg(
         hf_hub_id='timm/',
         interpolation='bicubic'),
-    'selecsls84.untrained': _cfg(
+    'SelecSls84.untrained': _cfg(
         interpolation='bicubic'),
 })
 
 
 @register_model
-def selecsls42(pretrained=False, **kwargs):
-    """Constructs a SelecSLS42 model.
+def SelecSls42(pretrained=False, **kwargs) -> SelecSls:
+    """Constructs a SelecSls42 model.
     """
-    return _create_selecsls('selecsls42', pretrained, **kwargs)
+    return _create_SelecSls('SelecSls42', pretrained, **kwargs)
 
 
 @register_model
-def selecsls42b(pretrained=False, **kwargs):
-    """Constructs a SelecSLS42_B model.
+def SelecSls42b(pretrained=False, **kwargs) -> SelecSls:
+    """Constructs a SelecSls42_B model.
     """
-    return _create_selecsls('selecsls42b', pretrained, **kwargs)
+    return _create_SelecSls('SelecSls42b', pretrained, **kwargs)
 
 
 @register_model
-def selecsls60(pretrained=False, **kwargs):
-    """Constructs a SelecSLS60 model.
+def SelecSls60(pretrained=False, **kwargs) -> SelecSls:
+    """Constructs a SelecSls60 model.
     """
-    return _create_selecsls('selecsls60', pretrained, **kwargs)
+    return _create_SelecSls('SelecSls60', pretrained, **kwargs)
 
 
 @register_model
-def selecsls60b(pretrained=False, **kwargs):
-    """Constructs a SelecSLS60_B model.
+def SelecSls60b(pretrained=False, **kwargs) -> SelecSls:
+    """Constructs a SelecSls60_B model.
     """
-    return _create_selecsls('selecsls60b', pretrained, **kwargs)
+    return _create_SelecSls('SelecSls60b', pretrained, **kwargs)
 
 
 @register_model
-def selecsls84(pretrained=False, **kwargs):
-    """Constructs a SelecSLS84 model.
+def SelecSls84(pretrained=False, **kwargs) -> SelecSls:
+    """Constructs a SelecSls84 model.
     """
-    return _create_selecsls('selecsls84', pretrained, **kwargs)
+    return _create_SelecSls('SelecSls84', pretrained, **kwargs)
diff --git a/timm/models/senet.py b/timm/models/senet.py
index 4706baa73e..8b203c372c 100644
--- a/timm/models/senet.py
+++ b/timm/models/senet.py
@@ -402,42 +402,42 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def legacy_seresnet18(pretrained=False, **kwargs):
+def legacy_seresnet18(pretrained=False, **kwargs) -> SENet:
     model_args = dict(
         block=SEResNetBlock, layers=[2, 2, 2, 2], groups=1, reduction=16, **kwargs)
     return _create_senet('legacy_seresnet18', pretrained, **model_args)
 
 
 @register_model
-def legacy_seresnet34(pretrained=False, **kwargs):
+def legacy_seresnet34(pretrained=False, **kwargs) -> SENet:
     model_args = dict(
         block=SEResNetBlock, layers=[3, 4, 6, 3], groups=1, reduction=16, **kwargs)
     return _create_senet('legacy_seresnet34', pretrained, **model_args)
 
 
 @register_model
-def legacy_seresnet50(pretrained=False, **kwargs):
+def legacy_seresnet50(pretrained=False, **kwargs) -> SENet:
     model_args = dict(
         block=SEResNetBottleneck, layers=[3, 4, 6, 3], groups=1, reduction=16, **kwargs)
     return _create_senet('legacy_seresnet50', pretrained, **model_args)
 
 
 @register_model
-def legacy_seresnet101(pretrained=False, **kwargs):
+def legacy_seresnet101(pretrained=False, **kwargs) -> SENet:
     model_args = dict(
         block=SEResNetBottleneck, layers=[3, 4, 23, 3], groups=1, reduction=16, **kwargs)
     return _create_senet('legacy_seresnet101', pretrained, **model_args)
 
 
 @register_model
-def legacy_seresnet152(pretrained=False, **kwargs):
+def legacy_seresnet152(pretrained=False, **kwargs) -> SENet:
     model_args = dict(
         block=SEResNetBottleneck, layers=[3, 8, 36, 3], groups=1, reduction=16, **kwargs)
     return _create_senet('legacy_seresnet152', pretrained, **model_args)
 
 
 @register_model
-def legacy_senet154(pretrained=False, **kwargs):
+def legacy_senet154(pretrained=False, **kwargs) -> SENet:
     model_args = dict(
         block=SEBottleneck, layers=[3, 8, 36, 3], groups=64, reduction=16,
         downsample_kernel_size=3, downsample_padding=1,  inplanes=128, input_3x3=True, **kwargs)
@@ -445,21 +445,21 @@ def legacy_senet154(pretrained=False, **kwargs):
 
 
 @register_model
-def legacy_seresnext26_32x4d(pretrained=False, **kwargs):
+def legacy_seresnext26_32x4d(pretrained=False, **kwargs) -> SENet:
     model_args = dict(
         block=SEResNeXtBottleneck, layers=[2, 2, 2, 2], groups=32, reduction=16, **kwargs)
     return _create_senet('legacy_seresnext26_32x4d', pretrained, **model_args)
 
 
 @register_model
-def legacy_seresnext50_32x4d(pretrained=False, **kwargs):
+def legacy_seresnext50_32x4d(pretrained=False, **kwargs) -> SENet:
     model_args = dict(
         block=SEResNeXtBottleneck, layers=[3, 4, 6, 3], groups=32, reduction=16, **kwargs)
     return _create_senet('legacy_seresnext50_32x4d', pretrained, **model_args)
 
 
 @register_model
-def legacy_seresnext101_32x4d(pretrained=False, **kwargs):
+def legacy_seresnext101_32x4d(pretrained=False, **kwargs) -> SENet:
     model_args = dict(
         block=SEResNeXtBottleneck, layers=[3, 4, 23, 3], groups=32, reduction=16, **kwargs)
     return _create_senet('legacy_seresnext101_32x4d', pretrained, **model_args)
diff --git a/timm/models/sequencer.py b/timm/models/sequencer.py
index 2899d29ed2..cb0f15f385 100644
--- a/timm/models/sequencer.py
+++ b/timm/models/sequencer.py
@@ -490,7 +490,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def sequencer2d_s(pretrained=False, **kwargs):
+def sequencer2d_s(pretrained=False, **kwargs) -> Sequencer2d:
     model_args = dict(
         layers=[4, 3, 8, 3],
         patch_sizes=[7, 2, 1, 1],
@@ -507,7 +507,7 @@ def sequencer2d_s(pretrained=False, **kwargs):
 
 
 @register_model
-def sequencer2d_m(pretrained=False, **kwargs):
+def sequencer2d_m(pretrained=False, **kwargs) -> Sequencer2d:
     model_args = dict(
         layers=[4, 3, 14, 3],
         patch_sizes=[7, 2, 1, 1],
@@ -524,7 +524,7 @@ def sequencer2d_m(pretrained=False, **kwargs):
 
 
 @register_model
-def sequencer2d_l(pretrained=False, **kwargs):
+def sequencer2d_l(pretrained=False, **kwargs) -> Sequencer2d:
     model_args = dict(
         layers=[8, 8, 16, 4],
         patch_sizes=[7, 2, 1, 1],
diff --git a/timm/models/sknet.py b/timm/models/sknet.py
index a0b8d6bf8a..01565875cb 100644
--- a/timm/models/sknet.py
+++ b/timm/models/sknet.py
@@ -172,7 +172,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def skresnet18(pretrained=False, **kwargs):
+def skresnet18(pretrained=False, **kwargs) -> ResNet:
     """Constructs a Selective Kernel ResNet-18 model.
 
     Different from configs in Select Kernel paper or "Compounding the Performance Improvements..." this
@@ -186,7 +186,7 @@ def skresnet18(pretrained=False, **kwargs):
 
 
 @register_model
-def skresnet34(pretrained=False, **kwargs):
+def skresnet34(pretrained=False, **kwargs) -> ResNet:
     """Constructs a Selective Kernel ResNet-34 model.
 
     Different from configs in Select Kernel paper or "Compounding the Performance Improvements..." this
@@ -200,7 +200,7 @@ def skresnet34(pretrained=False, **kwargs):
 
 
 @register_model
-def skresnet50(pretrained=False, **kwargs):
+def skresnet50(pretrained=False, **kwargs) -> ResNet:
     """Constructs a Select Kernel ResNet-50 model.
 
     Different from configs in Select Kernel paper or "Compounding the Performance Improvements..." this
@@ -214,7 +214,7 @@ def skresnet50(pretrained=False, **kwargs):
 
 
 @register_model
-def skresnet50d(pretrained=False, **kwargs):
+def skresnet50d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a Select Kernel ResNet-50-D model.
 
     Different from configs in Select Kernel paper or "Compounding the Performance Improvements..." this
@@ -228,7 +228,7 @@ def skresnet50d(pretrained=False, **kwargs):
 
 
 @register_model
-def skresnext50_32x4d(pretrained=False, **kwargs):
+def skresnext50_32x4d(pretrained=False, **kwargs) -> ResNet:
     """Constructs a Select Kernel ResNeXt50-32x4d model. This should be equivalent to
     the SKNet-50 model in the Select Kernel Paper
     """
diff --git a/timm/models/swin_transformer.py b/timm/models/swin_transformer.py
index 5cc99f23de..2cca973673 100644
--- a/timm/models/swin_transformer.py
+++ b/timm/models/swin_transformer.py
@@ -711,7 +711,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def swin_tiny_patch4_window7_224(pretrained=False, **kwargs):
+def swin_tiny_patch4_window7_224(pretrained=False, **kwargs) -> SwinTransformer:
     """ Swin-T @ 224x224, trained ImageNet-1k
     """
     model_args = dict(patch_size=4, window_size=7, embed_dim=96, depths=(2, 2, 6, 2), num_heads=(3, 6, 12, 24))
@@ -720,7 +720,7 @@ def swin_tiny_patch4_window7_224(pretrained=False, **kwargs):
 
 
 @register_model
-def swin_small_patch4_window7_224(pretrained=False, **kwargs):
+def swin_small_patch4_window7_224(pretrained=False, **kwargs) -> SwinTransformer:
     """ Swin-S @ 224x224
     """
     model_args = dict(patch_size=4, window_size=7, embed_dim=96, depths=(2, 2, 18, 2), num_heads=(3, 6, 12, 24))
@@ -729,7 +729,7 @@ def swin_small_patch4_window7_224(pretrained=False, **kwargs):
 
 
 @register_model
-def swin_base_patch4_window7_224(pretrained=False, **kwargs):
+def swin_base_patch4_window7_224(pretrained=False, **kwargs) -> SwinTransformer:
     """ Swin-B @ 224x224
     """
     model_args = dict(patch_size=4, window_size=7, embed_dim=128, depths=(2, 2, 18, 2), num_heads=(4, 8, 16, 32))
@@ -738,7 +738,7 @@ def swin_base_patch4_window7_224(pretrained=False, **kwargs):
 
 
 @register_model
-def swin_base_patch4_window12_384(pretrained=False, **kwargs):
+def swin_base_patch4_window12_384(pretrained=False, **kwargs) -> SwinTransformer:
     """ Swin-B @ 384x384
     """
     model_args = dict(patch_size=4, window_size=12, embed_dim=128, depths=(2, 2, 18, 2), num_heads=(4, 8, 16, 32))
@@ -747,7 +747,7 @@ def swin_base_patch4_window12_384(pretrained=False, **kwargs):
 
 
 @register_model
-def swin_large_patch4_window7_224(pretrained=False, **kwargs):
+def swin_large_patch4_window7_224(pretrained=False, **kwargs) -> SwinTransformer:
     """ Swin-L @ 224x224
     """
     model_args = dict(patch_size=4, window_size=7, embed_dim=192, depths=(2, 2, 18, 2), num_heads=(6, 12, 24, 48))
@@ -756,7 +756,7 @@ def swin_large_patch4_window7_224(pretrained=False, **kwargs):
 
 
 @register_model
-def swin_large_patch4_window12_384(pretrained=False, **kwargs):
+def swin_large_patch4_window12_384(pretrained=False, **kwargs) -> SwinTransformer:
     """ Swin-L @ 384x384
     """
     model_args = dict(patch_size=4, window_size=12, embed_dim=192, depths=(2, 2, 18, 2), num_heads=(6, 12, 24, 48))
@@ -765,7 +765,7 @@ def swin_large_patch4_window12_384(pretrained=False, **kwargs):
 
 
 @register_model
-def swin_s3_tiny_224(pretrained=False, **kwargs):
+def swin_s3_tiny_224(pretrained=False, **kwargs) -> SwinTransformer:
     """ Swin-S3-T @ 224x224, https://arxiv.org/abs/2111.14725
     """
     model_args = dict(
@@ -774,7 +774,7 @@ def swin_s3_tiny_224(pretrained=False, **kwargs):
 
 
 @register_model
-def swin_s3_small_224(pretrained=False, **kwargs):
+def swin_s3_small_224(pretrained=False, **kwargs) -> SwinTransformer:
     """ Swin-S3-S @ 224x224, https://arxiv.org/abs/2111.14725
     """
     model_args = dict(
@@ -783,7 +783,7 @@ def swin_s3_small_224(pretrained=False, **kwargs):
 
 
 @register_model
-def swin_s3_base_224(pretrained=False, **kwargs):
+def swin_s3_base_224(pretrained=False, **kwargs) -> SwinTransformer:
     """ Swin-S3-B @ 224x224, https://arxiv.org/abs/2111.14725
     """
     model_args = dict(
diff --git a/timm/models/swin_transformer_v2.py b/timm/models/swin_transformer_v2.py
index 861676e95e..dba74a9a38 100644
--- a/timm/models/swin_transformer_v2.py
+++ b/timm/models/swin_transformer_v2.py
@@ -711,7 +711,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def swinv2_tiny_window16_256(pretrained=False, **kwargs):
+def swinv2_tiny_window16_256(pretrained=False, **kwargs) -> SwinTransformerV2:
     """
     """
     model_args = dict(window_size=16, embed_dim=96, depths=(2, 2, 6, 2), num_heads=(3, 6, 12, 24))
@@ -720,7 +720,7 @@ def swinv2_tiny_window16_256(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_tiny_window8_256(pretrained=False, **kwargs):
+def swinv2_tiny_window8_256(pretrained=False, **kwargs) -> SwinTransformerV2:
     """
     """
     model_args = dict(window_size=8, embed_dim=96, depths=(2, 2, 6, 2), num_heads=(3, 6, 12, 24))
@@ -729,7 +729,7 @@ def swinv2_tiny_window8_256(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_small_window16_256(pretrained=False, **kwargs):
+def swinv2_small_window16_256(pretrained=False, **kwargs) -> SwinTransformerV2:
     """
     """
     model_args = dict(window_size=16, embed_dim=96, depths=(2, 2, 18, 2), num_heads=(3, 6, 12, 24))
@@ -738,7 +738,7 @@ def swinv2_small_window16_256(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_small_window8_256(pretrained=False, **kwargs):
+def swinv2_small_window8_256(pretrained=False, **kwargs) -> SwinTransformerV2:
     """
     """
     model_args = dict(window_size=8, embed_dim=96, depths=(2, 2, 18, 2), num_heads=(3, 6, 12, 24))
@@ -747,7 +747,7 @@ def swinv2_small_window8_256(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_base_window16_256(pretrained=False, **kwargs):
+def swinv2_base_window16_256(pretrained=False, **kwargs) -> SwinTransformerV2:
     """
     """
     model_args = dict(window_size=16, embed_dim=128, depths=(2, 2, 18, 2), num_heads=(4, 8, 16, 32))
@@ -756,7 +756,7 @@ def swinv2_base_window16_256(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_base_window8_256(pretrained=False, **kwargs):
+def swinv2_base_window8_256(pretrained=False, **kwargs) -> SwinTransformerV2:
     """
     """
     model_args = dict(window_size=8, embed_dim=128, depths=(2, 2, 18, 2), num_heads=(4, 8, 16, 32))
@@ -765,7 +765,7 @@ def swinv2_base_window8_256(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_base_window12_192(pretrained=False, **kwargs):
+def swinv2_base_window12_192(pretrained=False, **kwargs) -> SwinTransformerV2:
     """
     """
     model_args = dict(window_size=12, embed_dim=128, depths=(2, 2, 18, 2), num_heads=(4, 8, 16, 32))
@@ -774,7 +774,7 @@ def swinv2_base_window12_192(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_base_window12to16_192to256(pretrained=False, **kwargs):
+def swinv2_base_window12to16_192to256(pretrained=False, **kwargs) -> SwinTransformerV2:
     """
     """
     model_args = dict(
@@ -785,7 +785,7 @@ def swinv2_base_window12to16_192to256(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_base_window12to24_192to384(pretrained=False, **kwargs):
+def swinv2_base_window12to24_192to384(pretrained=False, **kwargs) -> SwinTransformerV2:
     """
     """
     model_args = dict(
@@ -796,7 +796,7 @@ def swinv2_base_window12to24_192to384(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_large_window12_192(pretrained=False, **kwargs):
+def swinv2_large_window12_192(pretrained=False, **kwargs) -> SwinTransformerV2:
     """
     """
     model_args = dict(window_size=12, embed_dim=192, depths=(2, 2, 18, 2), num_heads=(6, 12, 24, 48))
@@ -805,7 +805,7 @@ def swinv2_large_window12_192(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_large_window12to16_192to256(pretrained=False, **kwargs):
+def swinv2_large_window12to16_192to256(pretrained=False, **kwargs) -> SwinTransformerV2:
     """
     """
     model_args = dict(
@@ -816,7 +816,7 @@ def swinv2_large_window12to16_192to256(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_large_window12to24_192to384(pretrained=False, **kwargs):
+def swinv2_large_window12to24_192to384(pretrained=False, **kwargs) -> SwinTransformerV2:
     """
     """
     model_args = dict(
diff --git a/timm/models/swin_transformer_v2_cr.py b/timm/models/swin_transformer_v2_cr.py
index 2b7dbf16b4..cb26dde0e9 100644
--- a/timm/models/swin_transformer_v2_cr.py
+++ b/timm/models/swin_transformer_v2_cr.py
@@ -838,7 +838,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def swinv2_cr_tiny_384(pretrained=False, **kwargs):
+def swinv2_cr_tiny_384(pretrained=False, **kwargs) -> SwinTransformerV2Cr:
     """Swin-T V2 CR @ 384x384, trained ImageNet-1k"""
     model_args = dict(
         embed_dim=96,
@@ -849,7 +849,7 @@ def swinv2_cr_tiny_384(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_cr_tiny_224(pretrained=False, **kwargs):
+def swinv2_cr_tiny_224(pretrained=False, **kwargs) -> SwinTransformerV2Cr:
     """Swin-T V2 CR @ 224x224, trained ImageNet-1k"""
     model_args = dict(
         embed_dim=96,
@@ -860,7 +860,7 @@ def swinv2_cr_tiny_224(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_cr_tiny_ns_224(pretrained=False, **kwargs):
+def swinv2_cr_tiny_ns_224(pretrained=False, **kwargs) -> SwinTransformerV2Cr:
     """Swin-T V2 CR @ 224x224, trained ImageNet-1k w/ extra stage norms.
     ** Experimental, may make default if results are improved. **
     """
@@ -874,7 +874,7 @@ def swinv2_cr_tiny_ns_224(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_cr_small_384(pretrained=False, **kwargs):
+def swinv2_cr_small_384(pretrained=False, **kwargs) -> SwinTransformerV2Cr:
     """Swin-S V2 CR @ 384x384, trained ImageNet-1k"""
     model_args = dict(
         embed_dim=96,
@@ -885,7 +885,7 @@ def swinv2_cr_small_384(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_cr_small_224(pretrained=False, **kwargs):
+def swinv2_cr_small_224(pretrained=False, **kwargs) -> SwinTransformerV2Cr:
     """Swin-S V2 CR @ 224x224, trained ImageNet-1k"""
     model_args = dict(
         embed_dim=96,
@@ -896,7 +896,7 @@ def swinv2_cr_small_224(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_cr_small_ns_224(pretrained=False, **kwargs):
+def swinv2_cr_small_ns_224(pretrained=False, **kwargs) -> SwinTransformerV2Cr:
     """Swin-S V2 CR @ 224x224, trained ImageNet-1k"""
     model_args = dict(
         embed_dim=96,
@@ -908,7 +908,7 @@ def swinv2_cr_small_ns_224(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_cr_small_ns_256(pretrained=False, **kwargs):
+def swinv2_cr_small_ns_256(pretrained=False, **kwargs) -> SwinTransformerV2Cr:
     """Swin-S V2 CR @ 256x256, trained ImageNet-1k"""
     model_args = dict(
         embed_dim=96,
@@ -920,7 +920,7 @@ def swinv2_cr_small_ns_256(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_cr_base_384(pretrained=False, **kwargs):
+def swinv2_cr_base_384(pretrained=False, **kwargs) -> SwinTransformerV2Cr:
     """Swin-B V2 CR @ 384x384, trained ImageNet-1k"""
     model_args = dict(
         embed_dim=128,
@@ -931,7 +931,7 @@ def swinv2_cr_base_384(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_cr_base_224(pretrained=False, **kwargs):
+def swinv2_cr_base_224(pretrained=False, **kwargs) -> SwinTransformerV2Cr:
     """Swin-B V2 CR @ 224x224, trained ImageNet-1k"""
     model_args = dict(
         embed_dim=128,
@@ -942,7 +942,7 @@ def swinv2_cr_base_224(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_cr_base_ns_224(pretrained=False, **kwargs):
+def swinv2_cr_base_ns_224(pretrained=False, **kwargs) -> SwinTransformerV2Cr:
     """Swin-B V2 CR @ 224x224, trained ImageNet-1k"""
     model_args = dict(
         embed_dim=128,
@@ -954,7 +954,7 @@ def swinv2_cr_base_ns_224(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_cr_large_384(pretrained=False, **kwargs):
+def swinv2_cr_large_384(pretrained=False, **kwargs) -> SwinTransformerV2Cr:
     """Swin-L V2 CR @ 384x384, trained ImageNet-1k"""
     model_args = dict(
         embed_dim=192,
@@ -965,7 +965,7 @@ def swinv2_cr_large_384(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_cr_large_224(pretrained=False, **kwargs):
+def swinv2_cr_large_224(pretrained=False, **kwargs) -> SwinTransformerV2Cr:
     """Swin-L V2 CR @ 224x224, trained ImageNet-1k"""
     model_args = dict(
         embed_dim=192,
@@ -976,7 +976,7 @@ def swinv2_cr_large_224(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_cr_huge_384(pretrained=False, **kwargs):
+def swinv2_cr_huge_384(pretrained=False, **kwargs) -> SwinTransformerV2Cr:
     """Swin-H V2 CR @ 384x384, trained ImageNet-1k"""
     model_args = dict(
         embed_dim=352,
@@ -988,7 +988,7 @@ def swinv2_cr_huge_384(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_cr_huge_224(pretrained=False, **kwargs):
+def swinv2_cr_huge_224(pretrained=False, **kwargs) -> SwinTransformerV2Cr:
     """Swin-H V2 CR @ 224x224, trained ImageNet-1k"""
     model_args = dict(
         embed_dim=352,
@@ -1000,7 +1000,7 @@ def swinv2_cr_huge_224(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_cr_giant_384(pretrained=False, **kwargs):
+def swinv2_cr_giant_384(pretrained=False, **kwargs) -> SwinTransformerV2Cr:
     """Swin-G V2 CR @ 384x384, trained ImageNet-1k"""
     model_args = dict(
         embed_dim=512,
@@ -1012,7 +1012,7 @@ def swinv2_cr_giant_384(pretrained=False, **kwargs):
 
 
 @register_model
-def swinv2_cr_giant_224(pretrained=False, **kwargs):
+def swinv2_cr_giant_224(pretrained=False, **kwargs) -> SwinTransformerV2Cr:
     """Swin-G V2 CR @ 224x224, trained ImageNet-1k"""
     model_args = dict(
         embed_dim=512,
diff --git a/timm/models/tnt.py b/timm/models/tnt.py
index c99d277581..c359018764 100644
--- a/timm/models/tnt.py
+++ b/timm/models/tnt.py
@@ -355,7 +355,7 @@ def _create_tnt(variant, pretrained=False, **kwargs):
 
 
 @register_model
-def tnt_s_patch16_224(pretrained=False, **kwargs):
+def tnt_s_patch16_224(pretrained=False, **kwargs) -> TNT:
     model_cfg = dict(
         patch_size=16, embed_dim=384, inner_dim=24, depth=12, num_heads_outer=6,
         qkv_bias=False)
@@ -364,7 +364,7 @@ def tnt_s_patch16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def tnt_b_patch16_224(pretrained=False, **kwargs):
+def tnt_b_patch16_224(pretrained=False, **kwargs) -> TNT:
     model_cfg = dict(
         patch_size=16, embed_dim=640, inner_dim=40, depth=12, num_heads_outer=10,
         qkv_bias=False)
diff --git a/timm/models/tresnet.py b/timm/models/tresnet.py
index 4cbba636e4..57e4fb6a18 100644
--- a/timm/models/tresnet.py
+++ b/timm/models/tresnet.py
@@ -324,25 +324,25 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def tresnet_m(pretrained=False, **kwargs):
+def tresnet_m(pretrained=False, **kwargs) -> TResNet:
     model_kwargs = dict(layers=[3, 4, 11, 3], **kwargs)
     return _create_tresnet('tresnet_m', pretrained=pretrained, **model_kwargs)
 
 
 @register_model
-def tresnet_l(pretrained=False, **kwargs):
+def tresnet_l(pretrained=False, **kwargs) -> TResNet:
     model_kwargs = dict(layers=[4, 5, 18, 3], width_factor=1.2, **kwargs)
     return _create_tresnet('tresnet_l', pretrained=pretrained, **model_kwargs)
 
 
 @register_model
-def tresnet_xl(pretrained=False, **kwargs):
+def tresnet_xl(pretrained=False, **kwargs) -> TResNet:
     model_kwargs = dict(layers=[4, 5, 24, 3], width_factor=1.3, **kwargs)
     return _create_tresnet('tresnet_xl', pretrained=pretrained, **model_kwargs)
 
 
 @register_model
-def tresnet_v2_l(pretrained=False, **kwargs):
+def tresnet_v2_l(pretrained=False, **kwargs) -> TResNet:
     model_kwargs = dict(layers=[3, 4, 23, 3], width_factor=1.0, v2=True, **kwargs)
     return _create_tresnet('tresnet_v2_l', pretrained=pretrained, **model_kwargs)
 
diff --git a/timm/models/twins.py b/timm/models/twins.py
index dda2a5d1c7..b96a0234d0 100644
--- a/timm/models/twins.py
+++ b/timm/models/twins.py
@@ -458,7 +458,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def twins_pcpvt_small(pretrained=False, **kwargs):
+def twins_pcpvt_small(pretrained=False, **kwargs) -> Twins:
     model_args = dict(
         patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4],
         depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1])
@@ -466,7 +466,7 @@ def twins_pcpvt_small(pretrained=False, **kwargs):
 
 
 @register_model
-def twins_pcpvt_base(pretrained=False, **kwargs):
+def twins_pcpvt_base(pretrained=False, **kwargs) -> Twins:
     model_args = dict(
         patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4],
         depths=[3, 4, 18, 3], sr_ratios=[8, 4, 2, 1])
@@ -474,7 +474,7 @@ def twins_pcpvt_base(pretrained=False, **kwargs):
 
 
 @register_model
-def twins_pcpvt_large(pretrained=False, **kwargs):
+def twins_pcpvt_large(pretrained=False, **kwargs) -> Twins:
     model_args = dict(
         patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4],
         depths=[3, 8, 27, 3], sr_ratios=[8, 4, 2, 1])
@@ -482,7 +482,7 @@ def twins_pcpvt_large(pretrained=False, **kwargs):
 
 
 @register_model
-def twins_svt_small(pretrained=False, **kwargs):
+def twins_svt_small(pretrained=False, **kwargs) -> Twins:
     model_args = dict(
         patch_size=4, embed_dims=[64, 128, 256, 512], num_heads=[2, 4, 8, 16], mlp_ratios=[4, 4, 4, 4],
         depths=[2, 2, 10, 4], wss=[7, 7, 7, 7], sr_ratios=[8, 4, 2, 1])
@@ -490,7 +490,7 @@ def twins_svt_small(pretrained=False, **kwargs):
 
 
 @register_model
-def twins_svt_base(pretrained=False, **kwargs):
+def twins_svt_base(pretrained=False, **kwargs) -> Twins:
     model_args = dict(
         patch_size=4, embed_dims=[96, 192, 384, 768], num_heads=[3, 6, 12, 24], mlp_ratios=[4, 4, 4, 4],
         depths=[2, 2, 18, 2], wss=[7, 7, 7, 7], sr_ratios=[8, 4, 2, 1])
@@ -498,7 +498,7 @@ def twins_svt_base(pretrained=False, **kwargs):
 
 
 @register_model
-def twins_svt_large(pretrained=False, **kwargs):
+def twins_svt_large(pretrained=False, **kwargs) -> Twins:
     model_args = dict(
         patch_size=4, embed_dims=[128, 256, 512, 1024], num_heads=[4, 8, 16, 32], mlp_ratios=[4, 4, 4, 4],
         depths=[2, 2, 18, 2], wss=[7, 7, 7, 7], sr_ratios=[8, 4, 2, 1])
diff --git a/timm/models/visformer.py b/timm/models/visformer.py
index 15ae7bb9b8..4ae14adca2 100644
--- a/timm/models/visformer.py
+++ b/timm/models/visformer.py
@@ -463,7 +463,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def visformer_tiny(pretrained=False, **kwargs):
+def visformer_tiny(pretrained=False, **kwargs) -> Visformer:
     model_cfg = dict(
         init_channels=16, embed_dim=192, depth=(7, 4, 4), num_heads=3, mlp_ratio=4., group=8,
         attn_stage='011', spatial_conv='100', norm_layer=nn.BatchNorm2d, conv_init=True,
@@ -473,7 +473,7 @@ def visformer_tiny(pretrained=False, **kwargs):
 
 
 @register_model
-def visformer_small(pretrained=False, **kwargs):
+def visformer_small(pretrained=False, **kwargs) -> Visformer:
     model_cfg = dict(
         init_channels=32, embed_dim=384, depth=(7, 4, 4), num_heads=6, mlp_ratio=4., group=8,
         attn_stage='011', spatial_conv='100', norm_layer=nn.BatchNorm2d, conv_init=True,
diff --git a/timm/models/vision_transformer.py b/timm/models/vision_transformer.py
index 75697f972c..40e103a20e 100644
--- a/timm/models/vision_transformer.py
+++ b/timm/models/vision_transformer.py
@@ -1055,16 +1055,16 @@ def _cfg(url='', **kwargs):
         mean=(0., 0., 0.), std=(1., 1., 1.), crop_pct=0.875, interpolation='bilinear'),
 
     # Custom timm variants
-    'vit_base_patch16_rpn_224.in1k': _cfg(
+    'vit_base_patch16_rpn_224.sw_in1k': _cfg(
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tpu-weights/vit_base_patch16_rpn_224-sw-3b07e89d.pth',
         hf_hub_id='timm/'),
-    'vit_medium_patch16_gap_240.in12k': _cfg(
+    'vit_medium_patch16_gap_240.sw_in12k': _cfg(
         hf_hub_id='timm/',
         input_size=(3, 240, 240), crop_pct=0.95, num_classes=11821),
-    'vit_medium_patch16_gap_256.in12k_ft_in1k': _cfg(
+    'vit_medium_patch16_gap_256.sw_in12k_ft_in1k': _cfg(
         hf_hub_id='timm/',
         input_size=(3, 256, 256), crop_pct=0.95),
-    'vit_medium_patch16_gap_384.in12k_ft_in1k': _cfg(
+    'vit_medium_patch16_gap_384.sw_in12k_ft_in1k': _cfg(
         hf_hub_id='timm/',
         input_size=(3, 384, 384), crop_pct=0.95, crop_mode='squash'),
     'vit_base_patch16_gap_224': _cfg(),
@@ -1229,11 +1229,11 @@ def _cfg(url='', **kwargs):
         crop_pct=1.0, input_size=(3, 336, 336), num_classes=768),
 
     # experimental (may be removed)
-    'vit_base_patch32_plus_256': _cfg(url='', input_size=(3, 256, 256), crop_pct=0.95),
-    'vit_base_patch16_plus_240': _cfg(url='', input_size=(3, 240, 240), crop_pct=0.95),
-    'vit_small_patch16_36x1_224': _cfg(url=''),
-    'vit_small_patch16_18x2_224': _cfg(url=''),
-    'vit_base_patch16_18x2_224': _cfg(url=''),
+    'vit_base_patch32_plus_256.untrained': _cfg(url='', input_size=(3, 256, 256), crop_pct=0.95),
+    'vit_base_patch16_plus_240.untrained': _cfg(url='', input_size=(3, 240, 240), crop_pct=0.95),
+    'vit_small_patch16_36x1_224.untrained': _cfg(url=''),
+    'vit_small_patch16_18x2_224.untrained': _cfg(url=''),
+    'vit_base_patch16_18x2_224.untrained': _cfg(url=''),
 
     # EVA fine-tuned weights from MAE style MIM - EVA-CLIP target pretrain
     # https://github.com/baaivision/EVA/blob/7ecf2c0a370d97967e86d047d7af9188f78d2df3/eva/README.md#eva-l-learning-better-mim-representations-from-eva-clip
@@ -1332,14 +1332,16 @@ def _create_vision_transformer(variant, pretrained=False, **kwargs):
         _filter_fn = checkpoint_filter_fn
 
     return build_model_with_cfg(
-        VisionTransformer, variant, pretrained,
+        VisionTransformer,
+        variant,
+        pretrained,
         pretrained_filter_fn=_filter_fn,
         **kwargs,
     )
 
 
 @register_model
-def vit_tiny_patch16_224(pretrained=False, **kwargs):
+def vit_tiny_patch16_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Tiny (Vit-Ti/16)
     """
     model_args = dict(patch_size=16, embed_dim=192, depth=12, num_heads=3)
@@ -1348,7 +1350,7 @@ def vit_tiny_patch16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_tiny_patch16_384(pretrained=False, **kwargs):
+def vit_tiny_patch16_384(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Tiny (Vit-Ti/16) @ 384x384.
     """
     model_args = dict(patch_size=16, embed_dim=192, depth=12, num_heads=3)
@@ -1357,7 +1359,7 @@ def vit_tiny_patch16_384(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_small_patch32_224(pretrained=False, **kwargs):
+def vit_small_patch32_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Small (ViT-S/32)
     """
     model_args = dict(patch_size=32, embed_dim=384, depth=12, num_heads=6)
@@ -1366,7 +1368,7 @@ def vit_small_patch32_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_small_patch32_384(pretrained=False, **kwargs):
+def vit_small_patch32_384(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Small (ViT-S/32) at 384x384.
     """
     model_args = dict(patch_size=32, embed_dim=384, depth=12, num_heads=6)
@@ -1375,7 +1377,7 @@ def vit_small_patch32_384(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_small_patch16_224(pretrained=False, **kwargs):
+def vit_small_patch16_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Small (ViT-S/16)
     """
     model_args = dict(patch_size=16, embed_dim=384, depth=12, num_heads=6)
@@ -1384,7 +1386,7 @@ def vit_small_patch16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_small_patch16_384(pretrained=False, **kwargs):
+def vit_small_patch16_384(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Small (ViT-S/16)
     """
     model_args = dict(patch_size=16, embed_dim=384, depth=12, num_heads=6)
@@ -1393,7 +1395,7 @@ def vit_small_patch16_384(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_small_patch8_224(pretrained=False, **kwargs):
+def vit_small_patch8_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Small (ViT-S/8)
     """
     model_args = dict(patch_size=8, embed_dim=384, depth=12, num_heads=6)
@@ -1402,7 +1404,7 @@ def vit_small_patch8_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_base_patch32_224(pretrained=False, **kwargs):
+def vit_base_patch32_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Base (ViT-B/32) from original paper (https://arxiv.org/abs/2010.11929).
     ImageNet-1k weights fine-tuned from in21k, source https://github.com/google-research/vision_transformer.
     """
@@ -1412,7 +1414,7 @@ def vit_base_patch32_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_base_patch32_384(pretrained=False, **kwargs):
+def vit_base_patch32_384(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Base model (ViT-B/32) from original paper (https://arxiv.org/abs/2010.11929).
     ImageNet-1k weights fine-tuned from in21k @ 384x384, source https://github.com/google-research/vision_transformer.
     """
@@ -1422,7 +1424,7 @@ def vit_base_patch32_384(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_base_patch16_224(pretrained=False, **kwargs):
+def vit_base_patch16_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Base (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
     ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.
     """
@@ -1432,7 +1434,7 @@ def vit_base_patch16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_base_patch16_384(pretrained=False, **kwargs):
+def vit_base_patch16_384(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Base model (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
     ImageNet-1k weights fine-tuned from in21k @ 384x384, source https://github.com/google-research/vision_transformer.
     """
@@ -1442,7 +1444,7 @@ def vit_base_patch16_384(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_base_patch8_224(pretrained=False, **kwargs):
+def vit_base_patch8_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Base (ViT-B/8) from original paper (https://arxiv.org/abs/2010.11929).
     ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.
     """
@@ -1452,7 +1454,7 @@ def vit_base_patch8_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_large_patch32_224(pretrained=False, **kwargs):
+def vit_large_patch32_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Large model (ViT-L/32) from original paper (https://arxiv.org/abs/2010.11929). No pretrained weights.
     """
     model_args = dict(patch_size=32, embed_dim=1024, depth=24, num_heads=16)
@@ -1461,7 +1463,7 @@ def vit_large_patch32_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_large_patch32_384(pretrained=False, **kwargs):
+def vit_large_patch32_384(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Large model (ViT-L/32) from original paper (https://arxiv.org/abs/2010.11929).
     ImageNet-1k weights fine-tuned from in21k @ 384x384, source https://github.com/google-research/vision_transformer.
     """
@@ -1471,7 +1473,7 @@ def vit_large_patch32_384(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_large_patch16_224(pretrained=False, **kwargs):
+def vit_large_patch16_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).
     ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.
     """
@@ -1481,7 +1483,7 @@ def vit_large_patch16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_large_patch16_384(pretrained=False, **kwargs):
+def vit_large_patch16_384(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).
     ImageNet-1k weights fine-tuned from in21k @ 384x384, source https://github.com/google-research/vision_transformer.
     """
@@ -1491,7 +1493,7 @@ def vit_large_patch16_384(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_large_patch14_224(pretrained=False, **kwargs):
+def vit_large_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Large model (ViT-L/14)
     """
     model_args = dict(patch_size=14, embed_dim=1024, depth=24, num_heads=16)
@@ -1500,7 +1502,7 @@ def vit_large_patch14_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_huge_patch14_224(pretrained=False, **kwargs):
+def vit_huge_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Huge model (ViT-H/14) from original paper (https://arxiv.org/abs/2010.11929).
     """
     model_args = dict(patch_size=14, embed_dim=1280, depth=32, num_heads=16)
@@ -1509,7 +1511,7 @@ def vit_huge_patch14_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_giant_patch14_224(pretrained=False, **kwargs):
+def vit_giant_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Giant (little-g) model (ViT-g/14) from `Scaling Vision Transformers` - https://arxiv.org/abs/2106.04560
     """
     model_args = dict(patch_size=14, embed_dim=1408, mlp_ratio=48/11, depth=40, num_heads=16)
@@ -1518,7 +1520,7 @@ def vit_giant_patch14_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_gigantic_patch14_224(pretrained=False, **kwargs):
+def vit_gigantic_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Gigantic (big-G) model (ViT-G/14) from `Scaling Vision Transformers` - https://arxiv.org/abs/2106.04560
     """
     model_args = dict(patch_size=14, embed_dim=1664, mlp_ratio=64/13, depth=48, num_heads=16)
@@ -1528,7 +1530,7 @@ def vit_gigantic_patch14_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_base_patch16_224_miil(pretrained=False, **kwargs):
+def vit_base_patch16_224_miil(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Base (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
     Weights taken from: https://github.com/Alibaba-MIIL/ImageNet21K
     """
@@ -1539,7 +1541,7 @@ def vit_base_patch16_224_miil(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_medium_patch16_gap_240(pretrained=False, **kwargs):
+def vit_medium_patch16_gap_240(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Medium (ViT-M/16) w/o class token, w/ avg-pool @ 240x240
     """
     model_args = dict(
@@ -1551,7 +1553,7 @@ def vit_medium_patch16_gap_240(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_medium_patch16_gap_256(pretrained=False, **kwargs):
+def vit_medium_patch16_gap_256(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Medium (ViT-M/16) w/o class token, w/ avg-pool @ 256x256
     """
     model_args = dict(
@@ -1563,7 +1565,7 @@ def vit_medium_patch16_gap_256(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_medium_patch16_gap_384(pretrained=False, **kwargs):
+def vit_medium_patch16_gap_384(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Medium (ViT-M/16) w/o class token, w/ avg-pool @ 384x384
     """
     model_args = dict(
@@ -1575,7 +1577,7 @@ def vit_medium_patch16_gap_384(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_base_patch16_gap_224(pretrained=False, **kwargs):
+def vit_base_patch16_gap_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Base (ViT-B/16) w/o class token, w/ avg-pool @ 256x256
     """
     model_args = dict(
@@ -1586,7 +1588,7 @@ def vit_base_patch16_gap_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_base_patch32_clip_224(pretrained=False, **kwargs):
+def vit_base_patch32_clip_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-B/32 CLIP image tower @ 224x224
     """
     model_args = dict(
@@ -1597,7 +1599,7 @@ def vit_base_patch32_clip_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_base_patch32_clip_384(pretrained=False, **kwargs):
+def vit_base_patch32_clip_384(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-B/32 CLIP image tower @ 384x384
     """
     model_args = dict(
@@ -1608,7 +1610,7 @@ def vit_base_patch32_clip_384(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_base_patch32_clip_448(pretrained=False, **kwargs):
+def vit_base_patch32_clip_448(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-B/32 CLIP image tower @ 448x448
     """
     model_args = dict(
@@ -1619,7 +1621,7 @@ def vit_base_patch32_clip_448(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_base_patch16_clip_224(pretrained=False, **kwargs):
+def vit_base_patch16_clip_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-B/16 CLIP image tower
     """
     model_args = dict(patch_size=16, embed_dim=768, depth=12, num_heads=12, pre_norm=True, norm_layer=nn.LayerNorm)
@@ -1629,7 +1631,7 @@ def vit_base_patch16_clip_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_base_patch16_clip_384(pretrained=False, **kwargs):
+def vit_base_patch16_clip_384(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-B/16 CLIP image tower @ 384x384
     """
     model_args = dict(patch_size=16, embed_dim=768, depth=12, num_heads=12, pre_norm=True, norm_layer=nn.LayerNorm)
@@ -1639,7 +1641,7 @@ def vit_base_patch16_clip_384(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_large_patch14_clip_224(pretrained=False, **kwargs):
+def vit_large_patch14_clip_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Large model (ViT-L/14) CLIP image tower
     """
     model_args = dict(patch_size=14, embed_dim=1024, depth=24, num_heads=16, pre_norm=True, norm_layer=nn.LayerNorm)
@@ -1649,7 +1651,7 @@ def vit_large_patch14_clip_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_large_patch14_clip_336(pretrained=False, **kwargs):
+def vit_large_patch14_clip_336(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Large model (ViT-L/14) CLIP image tower @ 336x336
     """
     model_args = dict(patch_size=14, embed_dim=1024, depth=24, num_heads=16, pre_norm=True, norm_layer=nn.LayerNorm)
@@ -1659,7 +1661,7 @@ def vit_large_patch14_clip_336(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_huge_patch14_clip_224(pretrained=False, **kwargs):
+def vit_huge_patch14_clip_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Huge model (ViT-H/14) CLIP image tower.
     """
     model_args = dict(patch_size=14, embed_dim=1280, depth=32, num_heads=16, pre_norm=True, norm_layer=nn.LayerNorm)
@@ -1669,7 +1671,7 @@ def vit_huge_patch14_clip_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_huge_patch14_clip_336(pretrained=False, **kwargs):
+def vit_huge_patch14_clip_336(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Huge model (ViT-H/14) CLIP image tower @ 336x336
     """
     model_args = dict(patch_size=14, embed_dim=1280, depth=32, num_heads=16, pre_norm=True, norm_layer=nn.LayerNorm)
@@ -1679,7 +1681,7 @@ def vit_huge_patch14_clip_336(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_giant_patch14_clip_224(pretrained=False, **kwargs):
+def vit_giant_patch14_clip_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Giant (little-g) model (ViT-g/14) from `Scaling Vision Transformers` - https://arxiv.org/abs/2106.04560
     Pretrained weights from CLIP image tower.
     """
@@ -1691,7 +1693,7 @@ def vit_giant_patch14_clip_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_gigantic_patch14_clip_224(pretrained=False, **kwargs):
+def vit_gigantic_patch14_clip_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-bigG model (ViT-G/14) from `Scaling Vision Transformers` - https://arxiv.org/abs/2106.04560
     Pretrained weights from CLIP image tower.
     """
@@ -1704,7 +1706,7 @@ def vit_gigantic_patch14_clip_224(pretrained=False, **kwargs):
 # Experimental models below
 
 @register_model
-def vit_base_patch32_plus_256(pretrained=False, **kwargs):
+def vit_base_patch32_plus_256(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Base (ViT-B/32+)
     """
     model_args = dict(patch_size=32, embed_dim=896, depth=12, num_heads=14, init_values=1e-5)
@@ -1714,7 +1716,7 @@ def vit_base_patch32_plus_256(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_base_patch16_plus_240(pretrained=False, **kwargs):
+def vit_base_patch16_plus_240(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Base (ViT-B/16+)
     """
     model_args = dict(patch_size=16, embed_dim=896, depth=12, num_heads=14, init_values=1e-5)
@@ -1724,7 +1726,7 @@ def vit_base_patch16_plus_240(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_base_patch16_rpn_224(pretrained=False, **kwargs):
+def vit_base_patch16_rpn_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Base (ViT-B/16) w/ residual post-norm
     """
     model_args = dict(
@@ -1736,7 +1738,7 @@ def vit_base_patch16_rpn_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_small_patch16_36x1_224(pretrained=False, **kwargs):
+def vit_small_patch16_36x1_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Base w/ LayerScale + 36 x 1 (36 block serial) config. Experimental, may remove.
     Based on `Three things everyone should know about Vision Transformers` - https://arxiv.org/abs/2203.09795
     Paper focuses on 24x2 + 48x1 for 'Small' width but those are extremely slow.
@@ -1748,7 +1750,7 @@ def vit_small_patch16_36x1_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_small_patch16_18x2_224(pretrained=False, **kwargs):
+def vit_small_patch16_18x2_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Small w/ LayerScale + 18 x 2 (36 block parallel) config. Experimental, may remove.
     Based on `Three things everyone should know about Vision Transformers` - https://arxiv.org/abs/2203.09795
     Paper focuses on 24x2 + 48x1 for 'Small' width but those are extremely slow.
@@ -1761,7 +1763,7 @@ def vit_small_patch16_18x2_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_base_patch16_18x2_224(pretrained=False, **kwargs):
+def vit_base_patch16_18x2_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Base w/ LayerScale + 18 x 2 (36 block parallel) config. Experimental, may remove.
     Based on `Three things everyone should know about Vision Transformers` - https://arxiv.org/abs/2203.09795
     """
@@ -1773,7 +1775,7 @@ def vit_base_patch16_18x2_224(pretrained=False, **kwargs):
 
 
 @register_model
-def eva_large_patch14_196(pretrained=False, **kwargs):
+def eva_large_patch14_196(pretrained=False, **kwargs) -> VisionTransformer:
     """ EVA-large model https://arxiv.org/abs/2211.07636 /via MAE MIM pretrain"""
     model_args = dict(patch_size=14, embed_dim=1024, depth=24, num_heads=16, global_pool='avg')
     model = _create_vision_transformer(
@@ -1782,7 +1784,7 @@ def eva_large_patch14_196(pretrained=False, **kwargs):
 
 
 @register_model
-def eva_large_patch14_336(pretrained=False, **kwargs):
+def eva_large_patch14_336(pretrained=False, **kwargs) -> VisionTransformer:
     """ EVA-large model https://arxiv.org/abs/2211.07636 via MAE MIM pretrain"""
     model_args = dict(patch_size=14, embed_dim=1024, depth=24, num_heads=16, global_pool='avg')
     model = _create_vision_transformer('eva_large_patch14_336', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -1790,7 +1792,7 @@ def eva_large_patch14_336(pretrained=False, **kwargs):
 
 
 @register_model
-def flexivit_small(pretrained=False, **kwargs):
+def flexivit_small(pretrained=False, **kwargs) -> VisionTransformer:
     """ FlexiViT-Small
     """
     model_args = dict(patch_size=16, embed_dim=384, depth=12, num_heads=6, no_embed_class=True)
@@ -1799,7 +1801,7 @@ def flexivit_small(pretrained=False, **kwargs):
 
 
 @register_model
-def flexivit_base(pretrained=False, **kwargs):
+def flexivit_base(pretrained=False, **kwargs) -> VisionTransformer:
     """ FlexiViT-Base
     """
     model_args = dict(patch_size=16, embed_dim=768, depth=12, num_heads=12, no_embed_class=True)
@@ -1808,7 +1810,7 @@ def flexivit_base(pretrained=False, **kwargs):
 
 
 @register_model
-def flexivit_large(pretrained=False, **kwargs):
+def flexivit_large(pretrained=False, **kwargs) -> VisionTransformer:
     """ FlexiViT-Large
     """
     model_args = dict(patch_size=16, embed_dim=1024, depth=24, num_heads=16, no_embed_class=True)
@@ -1817,7 +1819,7 @@ def flexivit_large(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_base_patch16_xp_224(pretrained=False, **kwargs):
+def vit_base_patch16_xp_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Large model (ViT-L/14) w/ parallel blocks and qk norm enabled.
     """
     model_args = dict(
@@ -1830,7 +1832,7 @@ def vit_base_patch16_xp_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_large_patch14_xp_224(pretrained=False, **kwargs):
+def vit_large_patch14_xp_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Large model (ViT-L/14) w/ parallel blocks and qk norm enabled.
     """
     model_args = dict(
@@ -1843,7 +1845,7 @@ def vit_large_patch14_xp_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_huge_patch14_xp_224(pretrained=False, **kwargs):
+def vit_huge_patch14_xp_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Huge model (ViT-H/14) w/ parallel blocks and qk norm enabled.
     """
     model_args = dict(
diff --git a/timm/models/vision_transformer_hybrid.py b/timm/models/vision_transformer_hybrid.py
index 0108959a72..f41e90e89b 100644
--- a/timm/models/vision_transformer_hybrid.py
+++ b/timm/models/vision_transformer_hybrid.py
@@ -23,7 +23,7 @@
 from ._registry import generate_default_cfgs, register_model, register_model_deprecations
 from .resnet import resnet26d, resnet50d
 from .resnetv2 import ResNetV2, create_resnetv2_stem
-from .vision_transformer import _create_vision_transformer
+from .vision_transformer import _create_vision_transformer, VisionTransformer
 
 
 class HybridEmbed(nn.Module):
@@ -177,7 +177,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def vit_tiny_r_s16_p8_224(pretrained=False, **kwargs):
+def vit_tiny_r_s16_p8_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ R+ViT-Ti/S16 w/ 8x8 patch hybrid @ 224 x 224.
     """
     backbone = _resnetv2(layers=(), **kwargs)
@@ -188,7 +188,7 @@ def vit_tiny_r_s16_p8_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_tiny_r_s16_p8_384(pretrained=False, **kwargs):
+def vit_tiny_r_s16_p8_384(pretrained=False, **kwargs) -> VisionTransformer:
     """ R+ViT-Ti/S16 w/ 8x8 patch hybrid @ 384 x 384.
     """
     backbone = _resnetv2(layers=(), **kwargs)
@@ -199,7 +199,7 @@ def vit_tiny_r_s16_p8_384(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_small_r26_s32_224(pretrained=False, **kwargs):
+def vit_small_r26_s32_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ R26+ViT-S/S32 hybrid.
     """
     backbone = _resnetv2((2, 2, 2, 2), **kwargs)
@@ -210,7 +210,7 @@ def vit_small_r26_s32_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_small_r26_s32_384(pretrained=False, **kwargs):
+def vit_small_r26_s32_384(pretrained=False, **kwargs) -> VisionTransformer:
     """ R26+ViT-S/S32 hybrid.
     """
     backbone = _resnetv2((2, 2, 2, 2), **kwargs)
@@ -221,7 +221,7 @@ def vit_small_r26_s32_384(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_base_r26_s32_224(pretrained=False, **kwargs):
+def vit_base_r26_s32_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ R26+ViT-B/S32 hybrid.
     """
     backbone = _resnetv2((2, 2, 2, 2), **kwargs)
@@ -232,7 +232,7 @@ def vit_base_r26_s32_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_base_r50_s16_224(pretrained=False, **kwargs):
+def vit_base_r50_s16_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ R50+ViT-B/S16 hybrid from original paper (https://arxiv.org/abs/2010.11929).
     """
     backbone = _resnetv2((3, 4, 9), **kwargs)
@@ -243,7 +243,7 @@ def vit_base_r50_s16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_base_r50_s16_384(pretrained=False, **kwargs):
+def vit_base_r50_s16_384(pretrained=False, **kwargs) -> VisionTransformer:
     """ R50+ViT-B/16 hybrid from original paper (https://arxiv.org/abs/2010.11929).
     ImageNet-1k weights fine-tuned from in21k @ 384x384, source https://github.com/google-research/vision_transformer.
     """
@@ -255,7 +255,7 @@ def vit_base_r50_s16_384(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_large_r50_s32_224(pretrained=False, **kwargs):
+def vit_large_r50_s32_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ R50+ViT-L/S32 hybrid.
     """
     backbone = _resnetv2((3, 4, 6, 3), **kwargs)
@@ -266,7 +266,7 @@ def vit_large_r50_s32_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_large_r50_s32_384(pretrained=False, **kwargs):
+def vit_large_r50_s32_384(pretrained=False, **kwargs) -> VisionTransformer:
     """ R50+ViT-L/S32 hybrid.
     """
     backbone = _resnetv2((3, 4, 6, 3), **kwargs)
@@ -277,7 +277,7 @@ def vit_large_r50_s32_384(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_small_resnet26d_224(pretrained=False, **kwargs):
+def vit_small_resnet26d_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ Custom ViT small hybrid w/ ResNet26D stride 32. No pretrained weights.
     """
     backbone = resnet26d(pretrained=pretrained, in_chans=kwargs.get('in_chans', 3), features_only=True, out_indices=[4])
@@ -288,7 +288,7 @@ def vit_small_resnet26d_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_small_resnet50d_s16_224(pretrained=False, **kwargs):
+def vit_small_resnet50d_s16_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ Custom ViT small hybrid w/ ResNet50D 3-stages, stride 16. No pretrained weights.
     """
     backbone = resnet50d(pretrained=pretrained, in_chans=kwargs.get('in_chans', 3), features_only=True, out_indices=[3])
@@ -299,7 +299,7 @@ def vit_small_resnet50d_s16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_base_resnet26d_224(pretrained=False, **kwargs):
+def vit_base_resnet26d_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ Custom ViT base hybrid w/ ResNet26D stride 32. No pretrained weights.
     """
     backbone = resnet26d(pretrained=pretrained, in_chans=kwargs.get('in_chans', 3), features_only=True, out_indices=[4])
@@ -310,7 +310,7 @@ def vit_base_resnet26d_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_base_resnet50d_224(pretrained=False, **kwargs):
+def vit_base_resnet50d_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ Custom ViT base hybrid w/ ResNet50D stride 32. No pretrained weights.
     """
     backbone = resnet50d(pretrained=pretrained, in_chans=kwargs.get('in_chans', 3), features_only=True, out_indices=[4])
diff --git a/timm/models/vision_transformer_relpos.py b/timm/models/vision_transformer_relpos.py
index 0e962a62cf..ea428587c1 100644
--- a/timm/models/vision_transformer_relpos.py
+++ b/timm/models/vision_transformer_relpos.py
@@ -453,7 +453,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def vit_relpos_base_patch32_plus_rpn_256(pretrained=False, **kwargs):
+def vit_relpos_base_patch32_plus_rpn_256(pretrained=False, **kwargs) -> VisionTransformerRelPos:
     """ ViT-Base (ViT-B/32+) w/ relative log-coord position and residual post-norm, no class token
     """
     model_args = dict(patch_size=32, embed_dim=896, depth=12, num_heads=14, block_fn=ResPostRelPosBlock)
@@ -463,7 +463,7 @@ def vit_relpos_base_patch32_plus_rpn_256(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_relpos_base_patch16_plus_240(pretrained=False, **kwargs):
+def vit_relpos_base_patch16_plus_240(pretrained=False, **kwargs) -> VisionTransformerRelPos:
     """ ViT-Base (ViT-B/16+) w/ relative log-coord position, no class token
     """
     model_args = dict(patch_size=16, embed_dim=896, depth=12, num_heads=14)
@@ -473,7 +473,7 @@ def vit_relpos_base_patch16_plus_240(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_relpos_small_patch16_224(pretrained=False, **kwargs):
+def vit_relpos_small_patch16_224(pretrained=False, **kwargs) -> VisionTransformerRelPos:
     """ ViT-Base (ViT-B/16) w/ relative log-coord position, no class token
     """
     model_args = dict(patch_size=16, embed_dim=384, depth=12, num_heads=6, qkv_bias=False, fc_norm=True)
@@ -483,7 +483,7 @@ def vit_relpos_small_patch16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_relpos_medium_patch16_224(pretrained=False, **kwargs):
+def vit_relpos_medium_patch16_224(pretrained=False, **kwargs) -> VisionTransformerRelPos:
     """ ViT-Base (ViT-B/16) w/ relative log-coord position, no class token
     """
     model_args = dict(
@@ -494,7 +494,7 @@ def vit_relpos_medium_patch16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_relpos_base_patch16_224(pretrained=False, **kwargs):
+def vit_relpos_base_patch16_224(pretrained=False, **kwargs) -> VisionTransformerRelPos:
     """ ViT-Base (ViT-B/16) w/ relative log-coord position, no class token
     """
     model_args = dict(
@@ -505,7 +505,7 @@ def vit_relpos_base_patch16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_srelpos_small_patch16_224(pretrained=False, **kwargs):
+def vit_srelpos_small_patch16_224(pretrained=False, **kwargs) -> VisionTransformerRelPos:
     """ ViT-Base (ViT-B/16) w/ shared relative log-coord position, no class token
     """
     model_args = dict(
@@ -517,7 +517,7 @@ def vit_srelpos_small_patch16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_srelpos_medium_patch16_224(pretrained=False, **kwargs):
+def vit_srelpos_medium_patch16_224(pretrained=False, **kwargs) -> VisionTransformerRelPos:
     """ ViT-Base (ViT-B/16) w/ shared relative log-coord position, no class token
     """
     model_args = dict(
@@ -529,7 +529,7 @@ def vit_srelpos_medium_patch16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_relpos_medium_patch16_cls_224(pretrained=False, **kwargs):
+def vit_relpos_medium_patch16_cls_224(pretrained=False, **kwargs) -> VisionTransformerRelPos:
     """ ViT-Base (ViT-M/16) w/ relative log-coord position, class token present
     """
     model_args = dict(
@@ -541,7 +541,7 @@ def vit_relpos_medium_patch16_cls_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_relpos_base_patch16_cls_224(pretrained=False, **kwargs):
+def vit_relpos_base_patch16_cls_224(pretrained=False, **kwargs) -> VisionTransformerRelPos:
     """ ViT-Base (ViT-B/16) w/ relative log-coord position, class token present
     """
     model_args = dict(
@@ -552,7 +552,7 @@ def vit_relpos_base_patch16_cls_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_relpos_base_patch16_clsgap_224(pretrained=False, **kwargs):
+def vit_relpos_base_patch16_clsgap_224(pretrained=False, **kwargs) -> VisionTransformerRelPos:
     """ ViT-Base (ViT-B/16) w/ relative log-coord position, class token present
     NOTE this config is a bit of a mistake, class token was enabled but global avg-pool w/ fc-norm was not disabled
     Leaving here for comparisons w/ a future re-train as it performs quite well.
@@ -565,7 +565,7 @@ def vit_relpos_base_patch16_clsgap_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_relpos_small_patch16_rpn_224(pretrained=False, **kwargs):
+def vit_relpos_small_patch16_rpn_224(pretrained=False, **kwargs) -> VisionTransformerRelPos:
     """ ViT-Base (ViT-B/16) w/ relative log-coord position and residual post-norm, no class token
     """
     model_args = dict(
@@ -576,7 +576,7 @@ def vit_relpos_small_patch16_rpn_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_relpos_medium_patch16_rpn_224(pretrained=False, **kwargs):
+def vit_relpos_medium_patch16_rpn_224(pretrained=False, **kwargs) -> VisionTransformerRelPos:
     """ ViT-Base (ViT-B/16) w/ relative log-coord position and residual post-norm, no class token
     """
     model_args = dict(
@@ -587,7 +587,7 @@ def vit_relpos_medium_patch16_rpn_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_relpos_base_patch16_rpn_224(pretrained=False, **kwargs):
+def vit_relpos_base_patch16_rpn_224(pretrained=False, **kwargs) -> VisionTransformerRelPos:
     """ ViT-Base (ViT-B/16) w/ relative log-coord position and residual post-norm, no class token
     """
     model_args = dict(
diff --git a/timm/models/volo.py b/timm/models/volo.py
index f504b61d69..9401d84b9e 100644
--- a/timm/models/volo.py
+++ b/timm/models/volo.py
@@ -798,7 +798,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def volo_d1_224(pretrained=False, **kwargs):
+def volo_d1_224(pretrained=False, **kwargs) -> VOLO:
     """ VOLO-D1 model, Params: 27M """
     model_args = dict(layers=(4, 4, 8, 2), embed_dims=(192, 384, 384, 384), num_heads=(6, 12, 12, 12), **kwargs)
     model = _create_volo('volo_d1_224', pretrained=pretrained, **model_args)
@@ -806,7 +806,7 @@ def volo_d1_224(pretrained=False, **kwargs):
 
 
 @register_model
-def volo_d1_384(pretrained=False, **kwargs):
+def volo_d1_384(pretrained=False, **kwargs) -> VOLO:
     """ VOLO-D1 model, Params: 27M """
     model_args = dict(layers=(4, 4, 8, 2), embed_dims=(192, 384, 384, 384), num_heads=(6, 12, 12, 12), **kwargs)
     model = _create_volo('volo_d1_384', pretrained=pretrained, **model_args)
@@ -814,7 +814,7 @@ def volo_d1_384(pretrained=False, **kwargs):
 
 
 @register_model
-def volo_d2_224(pretrained=False, **kwargs):
+def volo_d2_224(pretrained=False, **kwargs) -> VOLO:
     """ VOLO-D2 model, Params: 59M """
     model_args = dict(layers=(6, 4, 10, 4), embed_dims=(256, 512, 512, 512), num_heads=(8, 16, 16, 16), **kwargs)
     model = _create_volo('volo_d2_224', pretrained=pretrained, **model_args)
@@ -822,7 +822,7 @@ def volo_d2_224(pretrained=False, **kwargs):
 
 
 @register_model
-def volo_d2_384(pretrained=False, **kwargs):
+def volo_d2_384(pretrained=False, **kwargs) -> VOLO:
     """ VOLO-D2 model, Params: 59M """
     model_args = dict(layers=(6, 4, 10, 4), embed_dims=(256, 512, 512, 512), num_heads=(8, 16, 16, 16), **kwargs)
     model = _create_volo('volo_d2_384', pretrained=pretrained, **model_args)
@@ -830,7 +830,7 @@ def volo_d2_384(pretrained=False, **kwargs):
 
 
 @register_model
-def volo_d3_224(pretrained=False, **kwargs):
+def volo_d3_224(pretrained=False, **kwargs) -> VOLO:
     """ VOLO-D3 model, Params: 86M """
     model_args = dict(layers=(8, 8, 16, 4), embed_dims=(256, 512, 512, 512), num_heads=(8, 16, 16, 16), **kwargs)
     model = _create_volo('volo_d3_224', pretrained=pretrained, **model_args)
@@ -838,7 +838,7 @@ def volo_d3_224(pretrained=False, **kwargs):
 
 
 @register_model
-def volo_d3_448(pretrained=False, **kwargs):
+def volo_d3_448(pretrained=False, **kwargs) -> VOLO:
     """ VOLO-D3 model, Params: 86M """
     model_args = dict(layers=(8, 8, 16, 4), embed_dims=(256, 512, 512, 512), num_heads=(8, 16, 16, 16), **kwargs)
     model = _create_volo('volo_d3_448', pretrained=pretrained, **model_args)
@@ -846,7 +846,7 @@ def volo_d3_448(pretrained=False, **kwargs):
 
 
 @register_model
-def volo_d4_224(pretrained=False, **kwargs):
+def volo_d4_224(pretrained=False, **kwargs) -> VOLO:
     """ VOLO-D4 model, Params: 193M """
     model_args = dict(layers=(8, 8, 16, 4), embed_dims=(384, 768, 768, 768), num_heads=(12, 16, 16, 16), **kwargs)
     model = _create_volo('volo_d4_224', pretrained=pretrained, **model_args)
@@ -854,7 +854,7 @@ def volo_d4_224(pretrained=False, **kwargs):
 
 
 @register_model
-def volo_d4_448(pretrained=False, **kwargs):
+def volo_d4_448(pretrained=False, **kwargs) -> VOLO:
     """ VOLO-D4 model, Params: 193M """
     model_args = dict(layers=(8, 8, 16, 4), embed_dims=(384, 768, 768, 768), num_heads=(12, 16, 16, 16), **kwargs)
     model = _create_volo('volo_d4_448', pretrained=pretrained, **model_args)
@@ -862,7 +862,7 @@ def volo_d4_448(pretrained=False, **kwargs):
 
 
 @register_model
-def volo_d5_224(pretrained=False, **kwargs):
+def volo_d5_224(pretrained=False, **kwargs) -> VOLO:
     """ VOLO-D5 model, Params: 296M
     stem_hidden_dim=128, the dim in patch embedding is 128 for VOLO-D5
     """
@@ -874,7 +874,7 @@ def volo_d5_224(pretrained=False, **kwargs):
 
 
 @register_model
-def volo_d5_448(pretrained=False, **kwargs):
+def volo_d5_448(pretrained=False, **kwargs) -> VOLO:
     """ VOLO-D5 model, Params: 296M
     stem_hidden_dim=128, the dim in patch embedding is 128 for VOLO-D5
     """
@@ -886,7 +886,7 @@ def volo_d5_448(pretrained=False, **kwargs):
 
 
 @register_model
-def volo_d5_512(pretrained=False, **kwargs):
+def volo_d5_512(pretrained=False, **kwargs) -> VOLO:
     """ VOLO-D5 model, Params: 296M
     stem_hidden_dim=128, the dim in patch embedding is 128 for VOLO-D5
     """
diff --git a/timm/models/vovnet.py b/timm/models/vovnet.py
index 60016918c3..f4a06065e8 100644
--- a/timm/models/vovnet.py
+++ b/timm/models/vovnet.py
@@ -417,54 +417,54 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def vovnet39a(pretrained=False, **kwargs):
+def vovnet39a(pretrained=False, **kwargs) -> VovNet:
     return _create_vovnet('vovnet39a', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def vovnet57a(pretrained=False, **kwargs):
+def vovnet57a(pretrained=False, **kwargs) -> VovNet:
     return _create_vovnet('vovnet57a', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def ese_vovnet19b_slim_dw(pretrained=False, **kwargs):
+def ese_vovnet19b_slim_dw(pretrained=False, **kwargs) -> VovNet:
     return _create_vovnet('ese_vovnet19b_slim_dw', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def ese_vovnet19b_dw(pretrained=False, **kwargs):
+def ese_vovnet19b_dw(pretrained=False, **kwargs) -> VovNet:
     return _create_vovnet('ese_vovnet19b_dw', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def ese_vovnet19b_slim(pretrained=False, **kwargs):
+def ese_vovnet19b_slim(pretrained=False, **kwargs) -> VovNet:
     return _create_vovnet('ese_vovnet19b_slim', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def ese_vovnet39b(pretrained=False, **kwargs):
+def ese_vovnet39b(pretrained=False, **kwargs) -> VovNet:
     return _create_vovnet('ese_vovnet39b', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def ese_vovnet57b(pretrained=False, **kwargs):
+def ese_vovnet57b(pretrained=False, **kwargs) -> VovNet:
     return _create_vovnet('ese_vovnet57b', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def ese_vovnet99b(pretrained=False, **kwargs):
+def ese_vovnet99b(pretrained=False, **kwargs) -> VovNet:
     return _create_vovnet('ese_vovnet99b', pretrained=pretrained, **kwargs)
 
 
 @register_model
-def eca_vovnet39b(pretrained=False, **kwargs):
+def eca_vovnet39b(pretrained=False, **kwargs) -> VovNet:
     return _create_vovnet('eca_vovnet39b', pretrained=pretrained, **kwargs)
 
 
 # Experimental Models
 
 @register_model
-def ese_vovnet39b_evos(pretrained=False, **kwargs):
+def ese_vovnet39b_evos(pretrained=False, **kwargs) -> VovNet:
     def norm_act_fn(num_features, **nkwargs):
         return create_norm_act_layer('evonorms0', num_features, jit=False, **nkwargs)
     return _create_vovnet('ese_vovnet39b_evos', pretrained=pretrained, norm_layer=norm_act_fn, **kwargs)
diff --git a/timm/models/xception.py b/timm/models/xception.py
index 14b6e4f120..041e2b9305 100644
--- a/timm/models/xception.py
+++ b/timm/models/xception.py
@@ -230,7 +230,6 @@ def _xception(variant, pretrained=False, **kwargs):
 
 default_cfgs = generate_default_cfgs({
     'legacy_xception.tf_in1k': {
-        'hf_hub_id': 'timm/',
         'url': 'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-cadene/xception-43020ad28.pth',
         'input_size': (3, 299, 299),
         'pool_size': (10, 10),
@@ -247,7 +246,7 @@ def _xception(variant, pretrained=False, **kwargs):
 
 
 @register_model
-def legacy_xception(pretrained=False, **kwargs):
+def legacy_xception(pretrained=False, **kwargs) -> Xception:
     return _xception('legacy_xception', pretrained=pretrained, **kwargs)
 
 
diff --git a/timm/models/xception_aligned.py b/timm/models/xception_aligned.py
index 119550cd9f..9f6560f155 100644
--- a/timm/models/xception_aligned.py
+++ b/timm/models/xception_aligned.py
@@ -320,7 +320,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def xception41(pretrained=False, **kwargs):
+def xception41(pretrained=False, **kwargs) -> XceptionAligned:
     """ Modified Aligned Xception-41
     """
     block_cfg = [
@@ -339,7 +339,7 @@ def xception41(pretrained=False, **kwargs):
 
 
 @register_model
-def xception65(pretrained=False, **kwargs):
+def xception65(pretrained=False, **kwargs) -> XceptionAligned:
     """ Modified Aligned Xception-65
     """
     block_cfg = [
@@ -358,7 +358,7 @@ def xception65(pretrained=False, **kwargs):
 
 
 @register_model
-def xception71(pretrained=False, **kwargs):
+def xception71(pretrained=False, **kwargs) -> XceptionAligned:
     """ Modified Aligned Xception-71
     """
     block_cfg = [
@@ -379,7 +379,7 @@ def xception71(pretrained=False, **kwargs):
 
 
 @register_model
-def xception41p(pretrained=False, **kwargs):
+def xception41p(pretrained=False, **kwargs) -> XceptionAligned:
     """ Modified Aligned Xception-41 w/ Pre-Act
     """
     block_cfg = [
@@ -398,7 +398,7 @@ def xception41p(pretrained=False, **kwargs):
 
 
 @register_model
-def xception65p(pretrained=False, **kwargs):
+def xception65p(pretrained=False, **kwargs) -> XceptionAligned:
     """ Modified Aligned Xception-65 w/ Pre-Act
     """
     block_cfg = [
diff --git a/timm/models/xcit.py b/timm/models/xcit.py
index fbabe12422..a4cf9e46cb 100644
--- a/timm/models/xcit.py
+++ b/timm/models/xcit.py
@@ -26,7 +26,7 @@
 from .cait import ClassAttn
 from .vision_transformer import Mlp
 
-__all__ = ['XCiT']  # model_registry will add each entrypoint fn to this
+__all__ = ['Xcit']  # model_registry will add each entrypoint fn to this
 
 
 @register_notrace_module  # reason: FX can't symbolically trace torch.arange in forward method
@@ -270,7 +270,7 @@ def forward(self, x, H: int, W: int):
         return x
 
 
-class XCiT(nn.Module):
+class Xcit(nn.Module):
     """
     Based on timm and DeiT code bases
     https://github.com/rwightman/pytorch-image-models/tree/master/timm
@@ -498,7 +498,7 @@ def checkpoint_filter_fn(state_dict, model):
 
 def _create_xcit(variant, pretrained=False, default_cfg=None, **kwargs):
     model = build_model_with_cfg(
-        XCiT,
+        Xcit,
         variant,
         pretrained,
         pretrained_filter_fn=checkpoint_filter_fn,
@@ -652,7 +652,7 @@ def _cfg(url='', **kwargs):
 
 
 @register_model
-def xcit_nano_12_p16_224(pretrained=False, **kwargs):
+def xcit_nano_12_p16_224(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=16, embed_dim=128, depth=12, num_heads=4, eta=1.0, tokens_norm=False)
     model = _create_xcit('xcit_nano_12_p16_224', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -660,7 +660,7 @@ def xcit_nano_12_p16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_nano_12_p16_384(pretrained=False, **kwargs):
+def xcit_nano_12_p16_384(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=16, embed_dim=128, depth=12, num_heads=4, eta=1.0, tokens_norm=False, img_size=384)
     model = _create_xcit('xcit_nano_12_p16_384', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -668,7 +668,7 @@ def xcit_nano_12_p16_384(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_tiny_12_p16_224(pretrained=False, **kwargs):
+def xcit_tiny_12_p16_224(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=16, embed_dim=192, depth=12, num_heads=4, eta=1.0, tokens_norm=True)
     model = _create_xcit('xcit_tiny_12_p16_224', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -676,7 +676,7 @@ def xcit_tiny_12_p16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_tiny_12_p16_384(pretrained=False, **kwargs):
+def xcit_tiny_12_p16_384(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=16, embed_dim=192, depth=12, num_heads=4, eta=1.0, tokens_norm=True)
     model = _create_xcit('xcit_tiny_12_p16_384', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -684,7 +684,7 @@ def xcit_tiny_12_p16_384(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_small_12_p16_224(pretrained=False, **kwargs):
+def xcit_small_12_p16_224(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=16, embed_dim=384, depth=12, num_heads=8, eta=1.0, tokens_norm=True)
     model = _create_xcit('xcit_small_12_p16_224', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -692,7 +692,7 @@ def xcit_small_12_p16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_small_12_p16_384(pretrained=False, **kwargs):
+def xcit_small_12_p16_384(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=16, embed_dim=384, depth=12, num_heads=8, eta=1.0, tokens_norm=True)
     model = _create_xcit('xcit_small_12_p16_384', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -700,7 +700,7 @@ def xcit_small_12_p16_384(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_tiny_24_p16_224(pretrained=False, **kwargs):
+def xcit_tiny_24_p16_224(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=16, embed_dim=192, depth=24, num_heads=4, eta=1e-5, tokens_norm=True)
     model = _create_xcit('xcit_tiny_24_p16_224', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -708,7 +708,7 @@ def xcit_tiny_24_p16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_tiny_24_p16_384(pretrained=False, **kwargs):
+def xcit_tiny_24_p16_384(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=16, embed_dim=192, depth=24, num_heads=4, eta=1e-5, tokens_norm=True)
     model = _create_xcit('xcit_tiny_24_p16_384', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -716,7 +716,7 @@ def xcit_tiny_24_p16_384(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_small_24_p16_224(pretrained=False, **kwargs):
+def xcit_small_24_p16_224(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=16, embed_dim=384, depth=24, num_heads=8, eta=1e-5, tokens_norm=True)
     model = _create_xcit('xcit_small_24_p16_224', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -724,7 +724,7 @@ def xcit_small_24_p16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_small_24_p16_384(pretrained=False, **kwargs):
+def xcit_small_24_p16_384(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=16, embed_dim=384, depth=24, num_heads=8, eta=1e-5, tokens_norm=True)
     model = _create_xcit('xcit_small_24_p16_384', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -732,7 +732,7 @@ def xcit_small_24_p16_384(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_medium_24_p16_224(pretrained=False, **kwargs):
+def xcit_medium_24_p16_224(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=16, embed_dim=512, depth=24, num_heads=8, eta=1e-5, tokens_norm=True)
     model = _create_xcit('xcit_medium_24_p16_224', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -740,7 +740,7 @@ def xcit_medium_24_p16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_medium_24_p16_384(pretrained=False, **kwargs):
+def xcit_medium_24_p16_384(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=16, embed_dim=512, depth=24, num_heads=8, eta=1e-5, tokens_norm=True)
     model = _create_xcit('xcit_medium_24_p16_384', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -748,7 +748,7 @@ def xcit_medium_24_p16_384(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_large_24_p16_224(pretrained=False, **kwargs):
+def xcit_large_24_p16_224(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=16, embed_dim=768, depth=24, num_heads=16, eta=1e-5, tokens_norm=True)
     model = _create_xcit('xcit_large_24_p16_224', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -756,7 +756,7 @@ def xcit_large_24_p16_224(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_large_24_p16_384(pretrained=False, **kwargs):
+def xcit_large_24_p16_384(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=16, embed_dim=768, depth=24, num_heads=16, eta=1e-5, tokens_norm=True)
     model = _create_xcit('xcit_large_24_p16_384', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -765,7 +765,7 @@ def xcit_large_24_p16_384(pretrained=False, **kwargs):
 
 # Patch size 8x8 models
 @register_model
-def xcit_nano_12_p8_224(pretrained=False, **kwargs):
+def xcit_nano_12_p8_224(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=8, embed_dim=128, depth=12, num_heads=4, eta=1.0, tokens_norm=False)
     model = _create_xcit('xcit_nano_12_p8_224', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -773,7 +773,7 @@ def xcit_nano_12_p8_224(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_nano_12_p8_384(pretrained=False, **kwargs):
+def xcit_nano_12_p8_384(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=8, embed_dim=128, depth=12, num_heads=4, eta=1.0, tokens_norm=False)
     model = _create_xcit('xcit_nano_12_p8_384', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -781,7 +781,7 @@ def xcit_nano_12_p8_384(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_tiny_12_p8_224(pretrained=False, **kwargs):
+def xcit_tiny_12_p8_224(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=8, embed_dim=192, depth=12, num_heads=4, eta=1.0, tokens_norm=True)
     model = _create_xcit('xcit_tiny_12_p8_224', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -789,7 +789,7 @@ def xcit_tiny_12_p8_224(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_tiny_12_p8_384(pretrained=False, **kwargs):
+def xcit_tiny_12_p8_384(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=8, embed_dim=192, depth=12, num_heads=4, eta=1.0, tokens_norm=True)
     model = _create_xcit('xcit_tiny_12_p8_384', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -797,7 +797,7 @@ def xcit_tiny_12_p8_384(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_small_12_p8_224(pretrained=False, **kwargs):
+def xcit_small_12_p8_224(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=8, embed_dim=384, depth=12, num_heads=8, eta=1.0, tokens_norm=True)
     model = _create_xcit('xcit_small_12_p8_224', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -805,7 +805,7 @@ def xcit_small_12_p8_224(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_small_12_p8_384(pretrained=False, **kwargs):
+def xcit_small_12_p8_384(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=8, embed_dim=384, depth=12, num_heads=8, eta=1.0, tokens_norm=True)
     model = _create_xcit('xcit_small_12_p8_384', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -813,7 +813,7 @@ def xcit_small_12_p8_384(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_tiny_24_p8_224(pretrained=False, **kwargs):
+def xcit_tiny_24_p8_224(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=8, embed_dim=192, depth=24, num_heads=4, eta=1e-5, tokens_norm=True)
     model = _create_xcit('xcit_tiny_24_p8_224', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -821,7 +821,7 @@ def xcit_tiny_24_p8_224(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_tiny_24_p8_384(pretrained=False, **kwargs):
+def xcit_tiny_24_p8_384(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=8, embed_dim=192, depth=24, num_heads=4, eta=1e-5, tokens_norm=True)
     model = _create_xcit('xcit_tiny_24_p8_384', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -829,7 +829,7 @@ def xcit_tiny_24_p8_384(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_small_24_p8_224(pretrained=False, **kwargs):
+def xcit_small_24_p8_224(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=8, embed_dim=384, depth=24, num_heads=8, eta=1e-5, tokens_norm=True)
     model = _create_xcit('xcit_small_24_p8_224', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -837,7 +837,7 @@ def xcit_small_24_p8_224(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_small_24_p8_384(pretrained=False, **kwargs):
+def xcit_small_24_p8_384(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=8, embed_dim=384, depth=24, num_heads=8, eta=1e-5, tokens_norm=True)
     model = _create_xcit('xcit_small_24_p8_384', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -845,7 +845,7 @@ def xcit_small_24_p8_384(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_medium_24_p8_224(pretrained=False, **kwargs):
+def xcit_medium_24_p8_224(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=8, embed_dim=512, depth=24, num_heads=8, eta=1e-5, tokens_norm=True)
     model = _create_xcit('xcit_medium_24_p8_224', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -853,7 +853,7 @@ def xcit_medium_24_p8_224(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_medium_24_p8_384(pretrained=False, **kwargs):
+def xcit_medium_24_p8_384(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=8, embed_dim=512, depth=24, num_heads=8, eta=1e-5, tokens_norm=True)
     model = _create_xcit('xcit_medium_24_p8_384', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -861,7 +861,7 @@ def xcit_medium_24_p8_384(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_large_24_p8_224(pretrained=False, **kwargs):
+def xcit_large_24_p8_224(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=8, embed_dim=768, depth=24, num_heads=16, eta=1e-5, tokens_norm=True)
     model = _create_xcit('xcit_large_24_p8_224', pretrained=pretrained, **dict(model_args, **kwargs))
@@ -869,7 +869,7 @@ def xcit_large_24_p8_224(pretrained=False, **kwargs):
 
 
 @register_model
-def xcit_large_24_p8_384(pretrained=False, **kwargs):
+def xcit_large_24_p8_384(pretrained=False, **kwargs) -> Xcit:
     model_args = dict(
         patch_size=8, embed_dim=768, depth=24, num_heads=16, eta=1e-5, tokens_norm=True)
     model = _create_xcit('xcit_large_24_p8_384', pretrained=pretrained, **dict(model_args, **kwargs))

From 8c6fccb879954a3715a73816f762aa3c36d2a636 Mon Sep 17 00:00:00 2001
From: Ross Wightman <rwightman@gmail.com>
Date: Mon, 8 May 2023 15:15:44 -0700
Subject: [PATCH 2/7] Allow passing state_dict directly via pretrained cfg
 mechanism as an override

---
 timm/models/_builder.py    | 17 +++++++++++++----
 timm/models/_pretrained.py | 11 ++++++-----
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/timm/models/_builder.py b/timm/models/_builder.py
index cbf057b29b..6170de3042 100644
--- a/timm/models/_builder.py
+++ b/timm/models/_builder.py
@@ -32,6 +32,7 @@ def _resolve_pretrained_source(pretrained_cfg):
     cfg_source = pretrained_cfg.get('source', '')
     pretrained_url = pretrained_cfg.get('url', None)
     pretrained_file = pretrained_cfg.get('file', None)
+    pretrained_sd = pretrained_cfg.get('state_dict', None)
     hf_hub_id = pretrained_cfg.get('hf_hub_id', None)
 
     # resolve where to load pretrained weights from
@@ -44,8 +45,13 @@ def _resolve_pretrained_source(pretrained_cfg):
         pretrained_loc = hf_hub_id
     else:
         # default source == timm or unspecified
-        if pretrained_file:
-            # file load override is the highest priority if set
+        if pretrained_sd:
+            # direct state_dict pass through is the highest priority
+            load_from = 'state_dict'
+            pretrained_loc = pretrained_sd
+            assert isinstance(pretrained_loc, dict)
+        elif pretrained_file:
+            # file load override is the second-highest priority if set
             load_from = 'file'
             pretrained_loc = pretrained_file
         else:
@@ -108,7 +114,7 @@ def load_custom_pretrained(
     if not load_from:
         _logger.warning("No pretrained weights exist for this model. Using random initialization.")
         return
-    if load_from == 'hf-hub':  # FIXME
+    if load_from == 'hf-hub':
         _logger.warning("Hugging Face hub not currently supported for custom load pretrained models.")
     elif load_from == 'url':
         pretrained_loc = download_cached_file(
@@ -150,7 +156,10 @@ def load_pretrained(
         return
 
     load_from, pretrained_loc = _resolve_pretrained_source(pretrained_cfg)
-    if load_from == 'file':
+    if load_from == 'state_dict':
+        _logger.info(f'Loading pretrained weights from state dict')
+        state_dict = pretrained_loc  # pretrained_loc is the actual state dict for this override
+    elif load_from == 'file':
         _logger.info(f'Loading pretrained weights from file ({pretrained_loc})')
         state_dict = load_state_dict(pretrained_loc)
     elif load_from == 'url':
diff --git a/timm/models/_pretrained.py b/timm/models/_pretrained.py
index 24de26e1e9..2938f8fe71 100644
--- a/timm/models/_pretrained.py
+++ b/timm/models/_pretrained.py
@@ -11,11 +11,12 @@
 class PretrainedCfg:
     """
     """
-    # weight locations
-    url: Optional[Union[str, Tuple[str, str]]] = None
-    file: Optional[str] = None
-    hf_hub_id: Optional[str] = None
-    hf_hub_filename: Optional[str] = None
+    # weight source locations
+    url: Optional[Union[str, Tuple[str, str]]] = None  # remote URL
+    file: Optional[str] = None  # local / shared filesystem path
+    state_dict: Optional[Dict[str, Any]] = None  # in-memory state dict
+    hf_hub_id: Optional[str] = None  # Hugging Face Hub model id ('organization/model')
+    hf_hub_filename: Optional[str] = None  # Hugging Face Hub filename (overrides default)
 
     source: Optional[str] = None  # source of cfg / weight location used (url, file, hf-hub)
     architecture: Optional[str] = None  # architecture variant can be set when not implicit

From 21e57c0b9e4540437d7c878935dfda80193dcb59 Mon Sep 17 00:00:00 2001
From: Ross Wightman <rwightman@gmail.com>
Date: Mon, 8 May 2023 17:03:51 -0700
Subject: [PATCH 3/7] Add missing beitv2 in1k -> in1k models

---
 timm/models/beit.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/timm/models/beit.py b/timm/models/beit.py
index d4fdfb8df9..fdfb5e59c3 100644
--- a/timm/models/beit.py
+++ b/timm/models/beit.py
@@ -477,6 +477,11 @@ def _cfg(url='', **kwargs):
         hf_hub_id='timm/',
         mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD
     ),
+    'beitv2_base_patch16_224.in1k_ft_in1k': _cfg(
+        url='https://conversationhub.blob.core.windows.net/beit-share-public/beitv2/beitv2_base_patch16_224_pt1k_ft1k.pth',
+        hf_hub_id='timm/',
+        mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD
+    ),
     'beitv2_base_patch16_224.in1k_ft_in22k': _cfg(
         url='https://conversationhub.blob.core.windows.net/beit-share-public/beitv2/beitv2_base_patch16_224_pt1k_ft21k.pth',
         hf_hub_id='timm/',
@@ -487,6 +492,11 @@ def _cfg(url='', **kwargs):
         hf_hub_id='timm/',
         crop_pct=0.95, mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD
     ),
+    'beitv2_large_patch16_224.in1k_ft_in1k': _cfg(
+        url='https://conversationhub.blob.core.windows.net/beit-share-public/beitv2/beitv2_large_patch16_224_pt1k_ft1k.pth',
+        hf_hub_id='timm/',
+        crop_pct=0.95, mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD
+    ),
     'beitv2_large_patch16_224.in1k_ft_in22k': _cfg(
         url='https://conversationhub.blob.core.windows.net/beit-share-public/beitv2/beitv2_large_patch16_224_pt1k_ft21k.pth',
         hf_hub_id='timm/',

From e3363a7159992ae27e84b4ed4dad21942018b650 Mon Sep 17 00:00:00 2001
From: Ross Wightman <rwightman@gmail.com>
Date: Tue, 9 May 2023 11:33:51 -0700
Subject: [PATCH 4/7] Support bitsandbytes optimizers in factory

---
 timm/optim/optim_factory.py | 55 +++++++++++++++++++++++++++++++++----
 1 file changed, 49 insertions(+), 6 deletions(-)

diff --git a/timm/optim/optim_factory.py b/timm/optim/optim_factory.py
index 10950210ed..2e3020119b 100644
--- a/timm/optim/optim_factory.py
+++ b/timm/optim/optim_factory.py
@@ -27,11 +27,6 @@
 from .rmsprop_tf import RMSpropTF
 from .sgdp import SGDP
 
-try:
-    from apex.optimizers import FusedNovoGrad, FusedAdam, FusedLAMB, FusedSGD
-    has_apex = True
-except ImportError:
-    has_apex = False
 
 _logger = logging.getLogger(__name__)
 
@@ -254,9 +249,23 @@ def create_optimizer_v2(
     opt_lower = opt.lower()
     opt_split = opt_lower.split('_')
     opt_lower = opt_split[-1]
-    if 'fused' in opt_lower:
+
+    if opt_lower.startswith('fused'):
+        try:
+            from apex.optimizers import FusedNovoGrad, FusedAdam, FusedLAMB, FusedSGD
+            has_apex = True
+        except ImportError:
+            has_apex = False
         assert has_apex and torch.cuda.is_available(), 'APEX and CUDA required for fused optimizers'
 
+    if opt_lower.startswith('bnb'):
+        try:
+            import bitsandbytes as bnb
+            has_bnb = True
+        except ImportError:
+            has_bnb = False
+        assert has_bnb and torch.cuda.is_available(), 'bitsandbytes and CUDA required for bnb optimizers'
+
     opt_args = dict(weight_decay=weight_decay, **kwargs)
 
     if lr is not None:
@@ -357,6 +366,40 @@ def create_optimizer_v2(
         opt_args.setdefault('betas', (0.95, 0.98))
         optimizer = FusedNovoGrad(parameters, **opt_args)
 
+    # bitsandbytes optimizers, require bitsandbytes to be installed
+    elif opt_lower == 'bnbsgd':
+        opt_args.pop('eps', None)
+        optimizer = bnb.optim.SGD(parameters, momentum=momentum, nesterov=True, **opt_args)
+    elif opt_lower == 'bnbsgd8bit':
+        opt_args.pop('eps', None)
+        optimizer = bnb.optim.SGD8bit(parameters, momentum=momentum, nesterov=True, **opt_args)
+    elif opt_lower == 'bnbmomentum':
+        opt_args.pop('eps', None)
+        optimizer = bnb.optim.SGD(parameters, momentum=momentum, **opt_args)
+    elif opt_lower == 'bnbmomentum8bit':
+        opt_args.pop('eps', None)
+        optimizer = bnb.optim.SGD8bit(parameters, momentum=momentum, **opt_args)
+    elif opt_lower == 'bnbadam':
+        optimizer = bnb.optim.Adam(parameters, **opt_args)
+    elif opt_lower == 'bnbadam8bit':
+        optimizer = bnb.optim.Adam8bit(parameters, **opt_args)
+    elif opt_lower == 'bnbadamw':
+        optimizer = bnb.optim.AdamW(parameters, **opt_args)
+    elif opt_lower == 'bnbadamw8bit':
+        optimizer = bnb.optim.AdamW8bit(parameters, **opt_args)
+    elif opt_lower == 'bnblamb':
+        optimizer = bnb.optim.LAMB(parameters, **opt_args)
+    elif opt_lower == 'bnblamb8bit':
+        optimizer = bnb.optim.LAMB8bit(parameters, **opt_args)
+    elif opt_lower == 'bnblars':
+        optimizer = bnb.optim.LARS(parameters, **opt_args)
+    elif opt_lower == 'bnblarsb8bit':
+        optimizer = bnb.optim.LAMB8bit(parameters, **opt_args)
+    elif opt_lower == 'bnblion':
+        optimizer = bnb.optim.Lion(parameters, **opt_args)
+    elif opt_lower == 'bnblion8bit':
+        optimizer = bnb.optim.Lion8bit(parameters, **opt_args)
+
     else:
         assert False and "Invalid optimizer"
         raise ValueError

From a01d8f86f40be4cd3ebbb4c58c46effb724e9f86 Mon Sep 17 00:00:00 2001
From: Ross Wightman <rwightman@gmail.com>
Date: Tue, 9 May 2023 17:59:22 -0700
Subject: [PATCH 5/7] Tweak DinoV2 add, add MAE ViT weights, add initial
 intermediate layer getter experiment

---
 timm/layers/__init__.py                  |   2 +-
 timm/layers/patch_embed.py               |  45 +++++++-
 timm/layers/pos_embed.py                 |  24 ++--
 timm/models/deit.py                      |  30 ++++-
 timm/models/vision_transformer.py        | 135 +++++++++++++++++------
 timm/models/vision_transformer_hybrid.py |  36 +++++-
 6 files changed, 222 insertions(+), 50 deletions(-)

diff --git a/timm/layers/__init__.py b/timm/layers/__init__.py
index 45f2e5410a..d55faccc4b 100644
--- a/timm/layers/__init__.py
+++ b/timm/layers/__init__.py
@@ -34,7 +34,7 @@
     SyncBatchNormAct, convert_sync_batchnorm, FrozenBatchNormAct2d, freeze_batch_norm_2d, unfreeze_batch_norm_2d
 from .padding import get_padding, get_same_padding, pad_same
 from .patch_dropout import PatchDropout
-from .patch_embed import PatchEmbed, resample_patch_embed
+from .patch_embed import PatchEmbed, PatchEmbedWithSize, resample_patch_embed
 from .pool2d_same import AvgPool2dSame, create_pool2d
 from .pos_embed import resample_abs_pos_embed
 from .pos_embed_rel import RelPosMlp, RelPosBias, RelPosBiasTf, gen_relative_position_index, gen_relative_log_coords
diff --git a/timm/layers/patch_embed.py b/timm/layers/patch_embed.py
index a4946531b2..6ca6a1a120 100644
--- a/timm/layers/patch_embed.py
+++ b/timm/layers/patch_embed.py
@@ -9,7 +9,7 @@
 Hacked together by / Copyright 2020 Ross Wightman
 """
 import logging
-from typing import List, Optional, Callable
+from typing import Callable, List, Optional, Tuple, Union
 
 import torch
 from torch import nn as nn
@@ -75,6 +75,49 @@ def forward(self, x):
         return x
 
 
+class PatchEmbedWithSize(PatchEmbed):
+    """ 2D Image to Patch Embedding
+    """
+    output_fmt: Format
+
+    def __init__(
+            self,
+            img_size: Optional[int] = 224,
+            patch_size: int = 16,
+            in_chans: int = 3,
+            embed_dim: int = 768,
+            norm_layer: Optional[Callable] = None,
+            flatten: bool = True,
+            output_fmt: Optional[str] = None,
+            bias: bool = True,
+    ):
+        super().__init__(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+            norm_layer=norm_layer,
+            flatten=flatten,
+            output_fmt=output_fmt,
+            bias=bias,
+        )
+
+    def forward(self, x) -> Tuple[torch.Tensor, List[int]]:
+        B, C, H, W = x.shape
+        if self.img_size is not None:
+            _assert(H % self.patch_size[0] == 0, f"Input image height ({H}) must be divisible by patch size ({self.patch_size[0]}).")
+            _assert(W % self.patch_size[1] == 0, f"Input image width ({W}) must be divisible by patch size ({self.patch_size[1]}).")
+
+        x = self.proj(x)
+        grid_size = x.shape[-2:]
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # NCHW -> NLC
+        elif self.output_fmt != Format.NCHW:
+            x = nchw_to(x, self.output_fmt)
+        x = self.norm(x)
+        return x, grid_size
+
+
 def resample_patch_embed(
         patch_embed,
         new_size: List[int],
diff --git a/timm/layers/pos_embed.py b/timm/layers/pos_embed.py
index d0e675210c..c3afce76cc 100644
--- a/timm/layers/pos_embed.py
+++ b/timm/layers/pos_embed.py
@@ -24,29 +24,31 @@ def resample_abs_pos_embed(
         verbose: bool = False,
 ):
     # sort out sizes, assume square if old size not provided
-    new_size = to_2tuple(new_size)
-    new_ntok = new_size[0] * new_size[1]
-    if not old_size:
-        old_size = int(math.sqrt(posemb.shape[1] - num_prefix_tokens))
-    old_size = to_2tuple(old_size)
-    if new_size == old_size:  # might not both be same container type
+    num_pos_tokens = posemb.shape[1]
+    num_new_tokens = new_size[0] * new_size[1] + num_prefix_tokens
+    if num_new_tokens == num_pos_tokens and new_size[0] == new_size[1]:
         return posemb
 
+    if not old_size:
+        hw = int(math.sqrt(num_pos_tokens - num_prefix_tokens))
+        old_size = hw, hw
+
     if num_prefix_tokens:
         posemb_prefix, posemb = posemb[:, :num_prefix_tokens], posemb[:, num_prefix_tokens:]
     else:
         posemb_prefix, posemb = None, posemb
 
     # do the interpolation
+    embed_dim = posemb.shape[-1]
     posemb = posemb.reshape(1, old_size[0], old_size[1], -1).permute(0, 3, 1, 2)
     posemb = F.interpolate(posemb, size=new_size, mode=interpolation, antialias=antialias)
-    posemb = posemb.permute(0, 2, 3, 1).reshape(1, new_ntok, -1)
-
-    if verbose:
-        _logger.info(f'Resized position embedding: {old_size} to {new_size}.')
+    posemb = posemb.permute(0, 2, 3, 1).reshape(1, -1, embed_dim)
 
     # add back extra (class, etc) prefix tokens
     if posemb_prefix is not None:
-        print(posemb_prefix.shape, posemb.shape)
         posemb = torch.cat([posemb_prefix, posemb], dim=1)
+
+    if not torch.jit.is_scripting() and verbose:
+        _logger.info(f'Resized position embedding: {old_size} to {new_size}.')
+
     return posemb
diff --git a/timm/models/deit.py b/timm/models/deit.py
index be52a977e6..650ab6796e 100644
--- a/timm/models/deit.py
+++ b/timm/models/deit.py
@@ -11,11 +11,13 @@
 # Copyright (c) 2015-present, Facebook, Inc.
 # All rights reserved.
 from functools import partial
+from typing import Sequence, Union
 
 import torch
 from torch import nn as nn
 
 from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+from timm.layers import resample_abs_pos_embed
 from timm.models.vision_transformer import VisionTransformer, trunc_normal_, checkpoint_filter_fn
 from ._builder import build_model_with_cfg
 from ._manipulate import checkpoint_seq
@@ -71,11 +73,37 @@ def reset_classifier(self, num_classes, global_pool=None):
     def set_distilled_training(self, enable=True):
         self.distilled_training = enable
 
+    def _intermediate_layers(
+            self,
+            x: torch.Tensor,
+            n: Union[int, Sequence] = 1,
+    ):
+        outputs, num_blocks = [], len(self.blocks)
+        take_indices = set(range(num_blocks - n, num_blocks) if isinstance(n, int) else n)
+
+        # forward pass
+        x = self.patch_embed(x)
+        x = torch.cat((
+            self.cls_token.expand(x.shape[0], -1, -1),
+            self.dist_token.expand(x.shape[0], -1, -1),
+            x),
+            dim=1)
+        x = self.pos_drop(x + self.pos_embed)
+        x = self.patch_drop(x)
+        x = self.norm_pre(x)
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if i in take_indices:
+                outputs.append(x)
+
+        return outputs
+
     def forward_features(self, x) -> torch.Tensor:
         x = self.patch_embed(x)
         x = torch.cat((
             self.cls_token.expand(x.shape[0], -1, -1),
-            self.dist_token.expand(x.shape[0], -1, -1), x),
+            self.dist_token.expand(x.shape[0], -1, -1),
+            x),
             dim=1)
         x = self.pos_drop(x + self.pos_embed)
         if self.grad_checkpointing and not torch.jit.is_scripting():
diff --git a/timm/models/vision_transformer.py b/timm/models/vision_transformer.py
index b38a168dab..a418f2b0e3 100644
--- a/timm/models/vision_transformer.py
+++ b/timm/models/vision_transformer.py
@@ -27,7 +27,7 @@
 import math
 from collections import OrderedDict
 from functools import partial
-from typing import Callable, List, Optional, Tuple, Union
+from typing import Callable, List, Optional, Sequence, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -125,7 +125,7 @@ def __init__(
             drop_path=0.,
             act_layer=nn.GELU,
             norm_layer=nn.LayerNorm,
-            ffn_layer=Mlp,
+            mlp_layer=Mlp,
     ):
         super().__init__()
         self.norm1 = norm_layer(dim)
@@ -142,7 +142,7 @@ def __init__(
         self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
 
         self.norm2 = norm_layer(dim)
-        self.mlp = ffn_layer(
+        self.mlp = mlp_layer(
             in_features=dim,
             hidden_features=int(dim * mlp_ratio),
             act_layer=act_layer,
@@ -172,7 +172,7 @@ def __init__(
             drop_path=0.,
             act_layer=nn.GELU,
             norm_layer=nn.LayerNorm,
-            ffn_layer=Mlp,
+            mlp_layer=Mlp,
     ):
         super().__init__()
         self.init_values = init_values
@@ -189,7 +189,7 @@ def __init__(
         self.norm1 = norm_layer(dim)
         self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
 
-        self.mlp = ffn_layer(
+        self.mlp = mlp_layer(
             in_features=dim,
             hidden_features=int(dim * mlp_ratio),
             act_layer=act_layer,
@@ -232,7 +232,7 @@ def __init__(
             drop_path=0.,
             act_layer=nn.GELU,
             norm_layer=nn.LayerNorm,
-            ffn_layer=None, # NOTE: not used
+            mlp_layer=None,  # NOTE: not used
     ):
         super().__init__()
         assert dim % num_heads == 0, 'dim should be divisible by num_heads'
@@ -326,7 +326,7 @@ def __init__(
             drop_path=0.,
             act_layer=nn.GELU,
             norm_layer=nn.LayerNorm,
-            ffn_layer=Mlp,
+            mlp_layer=Mlp,
     ):
         super().__init__()
         self.num_parallel = num_parallel
@@ -349,7 +349,7 @@ def __init__(
             ])))
             self.ffns.append(nn.Sequential(OrderedDict([
                 ('norm', norm_layer(dim)),
-                ('mlp', ffn_layer(
+                ('mlp', mlp_layer(
                     dim,
                     hidden_features=int(dim * mlp_ratio),
                     act_layer=act_layer,
@@ -413,7 +413,7 @@ def __init__(
             norm_layer: Optional[Callable] = None,
             act_layer: Optional[Callable] = None,
             block_fn: Callable = Block,
-            ffn_layer: Callable = Mlp,
+            mlp_layer: Callable = Mlp,
     ):
         """
         Args:
@@ -435,7 +435,7 @@ def __init__(
             attn_drop_rate: Attention dropout rate.
             drop_path_rate: Stochastic depth rate.
             weight_init: Weight initialization scheme.
-            embed_layer: Patch embedding layey.
+            embed_layer: Patch embedding layer.
             norm_layer: Normalization layer.
             act_layer: MLP activation layer.
             block_fn: Transformer block layer.
@@ -490,7 +490,7 @@ def __init__(
                 drop_path=dpr[i],
                 norm_layer=norm_layer,
                 act_layer=act_layer,
-                ffn_layer=ffn_layer,
+                mlp_layer=mlp_layer,
             )
             for i in range(depth)])
         self.norm = norm_layer(embed_dim) if not use_fc_norm else nn.Identity()
@@ -560,6 +560,55 @@ def _pos_embed(self, x):
             x = x + self.pos_embed
         return self.pos_drop(x)
 
+    def _intermediate_layers(
+            self,
+            x: torch.Tensor,
+            n: Union[int, Sequence] = 1,
+    ):
+        outputs, num_blocks = [], len(self.blocks)
+        take_indices = set(range(num_blocks - n, num_blocks) if isinstance(n, int) else n)
+
+        # forward pass
+        x = self.patch_embed(x)
+        x = self._pos_embed(x)
+        x = self.patch_drop(x)
+        x = self.norm_pre(x)
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if i in take_indices:
+                outputs.append(x)
+
+        return outputs
+
+    def get_intermediate_layers(
+            self,
+            x: torch.Tensor,
+            n: Union[int, Sequence] = 1,
+            reshape: bool = False,
+            return_class_token: bool = False,
+            norm: bool = False,
+    ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:
+        """ Intermediate layer accessor (NOTE: This is a WIP experiment).
+        Inspired by DINO / DINOv2 interface
+        """
+        # take last n blocks if n is an int, if in is a sequence, select by matching indices
+        outputs = self._intermediate_layers(x, n)
+        if norm:
+            outputs = [self.norm(out) for out in outputs]
+        class_tokens = [out[:, 0:self.num_prefix_tokens] for out in outputs]
+        outputs = [out[:, self.num_prefix_tokens:] for out in outputs]
+
+        if reshape:
+            grid_size = self.patch_embed.grid_size
+            outputs = [
+                out.reshape(x.shape[0], grid_size[0], grid_size[1], -1).permute(0, 3, 1, 2).contiguous()
+                for out in outputs
+            ]
+
+        if return_class_token:
+            return tuple(zip(outputs, class_tokens))
+        return tuple(outputs)
+
     def forward_features(self, x):
         x = self.patch_embed(x)
         x = self._pos_embed(x)
@@ -816,9 +865,7 @@ def _convert_openai_clip(state_dict, model):
 
 def _convert_dinov2(state_dict, model):
     import re
-
     out_dict = {}
-
     for k, v in state_dict.items():
         if k == "mask_token":
             continue
@@ -828,11 +875,10 @@ def _convert_dinov2(state_dict, model):
         elif re.match(r"blocks\.(\d+)\.mlp\.w3\.(?:weight|bias)", k):
             out_dict[k.replace("w3", "fc2")] = v
             continue
-    
         out_dict[k] = v
-
     return out_dict
 
+
 def checkpoint_filter_fn(
         state_dict,
         model,
@@ -1072,19 +1118,27 @@ def _cfg(url='', **kwargs):
         hf_hub_id='timm/',
         mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD, num_classes=0),
     
-    # DINOv2 pretrained - https://arxiv.org/abs/2304.07193 (no classifier head, for fine-tune only)
-    'vit_small_patch14_dinov2': _cfg(
+    # DINOv2 pretrained - https://arxiv.org/abs/2304.07193 (no classifier head, for fine-tune/features only)
+    'vit_small_patch14_dinov2.lvd142m': _cfg(
         url='https://dl.fbaipublicfiles.com/dinov2/dinov2_vits14/dinov2_vits14_pretrain.pth',
-        mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD, num_classes=0, input_size=(3, 518, 518)),
-    'vit_base_patch14_dinov2': _cfg(
+        license='cc-by-nc-4.0',
+        mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD, num_classes=0,
+        input_size=(3, 518, 518), crop_pct=1.0),
+    'vit_base_patch14_dinov2.lvd142m': _cfg(
         url='https://dl.fbaipublicfiles.com/dinov2/dinov2_vitb14/dinov2_vitb14_pretrain.pth',
-        mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD, num_classes=0, input_size=(3, 518, 518)),
-    'vit_large_patch14_dinov2': _cfg(
+        license='cc-by-nc-4.0',
+        mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD, num_classes=0,
+        input_size=(3, 518, 518), crop_pct=1.0),
+    'vit_large_patch14_dinov2.lvd142m': _cfg(
         url='https://dl.fbaipublicfiles.com/dinov2/dinov2_vitl14/dinov2_vitl14_pretrain.pth',
-        mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD, num_classes=0, input_size=(3, 518, 518)),
-    'vit_giant_patch14_dinov2': _cfg(
+        license='cc-by-nc-4.0',
+        mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD, num_classes=0,
+        input_size=(3, 518, 518), crop_pct=1.0),
+    'vit_giant_patch14_dinov2.lvd142m': _cfg(
         url='https://dl.fbaipublicfiles.com/dinov2/dinov2_vitg14/dinov2_vitg14_pretrain.pth',
-        mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD, num_classes=0, input_size=(3, 518, 518)),
+        license='cc-by-nc-4.0',
+        mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD, num_classes=0,
+        input_size=(3, 518, 518), crop_pct=1.0),
 
     # ViT ImageNet-21K-P pretraining by MILL
     'vit_base_patch16_224_miil.in21k': _cfg(
@@ -1359,6 +1413,22 @@ def _cfg(url='', **kwargs):
     'vit_base_patch16_xp_224.untrained': _cfg(url=''),
     'vit_large_patch14_xp_224.untrained': _cfg(url=''),
     'vit_huge_patch14_xp_224.untrained': _cfg(url=''),
+
+    'vit_base_patch16_224.mae': _cfg(
+        url='https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_base.pth',
+        #hf_hub_id='timm/',
+        license='cc-by-nc-4.0',
+        mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD, num_classes=0),
+    'vit_large_patch16_224.mae': _cfg(
+        url='https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_large.pth',
+        # hf_hub_id='timm/',
+        license='cc-by-nc-4.0',
+        mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD, num_classes=0),
+    'vit_huge_patch14_224.mae': _cfg(
+        url='https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_huge.pth',
+        # hf_hub_id='timm/',
+        license='cc-by-nc-4.0',
+        mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD, num_classes=0),
 })
 
 
@@ -1904,10 +1974,8 @@ def vit_small_patch14_dinov2(pretrained=False, **kwargs):
     """ ViT-S/14 for DINOv2
     """
     model_args = dict(
-        patch_size=14, embed_dim=384, depth=12, num_heads=6, 
-        init_values=1.0, img_size=518, 
+        patch_size=14, embed_dim=384, depth=12, num_heads=6, init_values=1.0, img_size=518,
     )
-    
     model = _create_vision_transformer(
         'vit_small_patch14_dinov2', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
@@ -1918,10 +1986,8 @@ def vit_base_patch14_dinov2(pretrained=False, **kwargs):
     """ ViT-B/14 for DINOv2
     """
     model_args = dict(
-        patch_size=14, embed_dim=768, depth=12, num_heads=12, 
-        init_values=1.0, img_size=518, 
+        patch_size=14, embed_dim=768, depth=12, num_heads=12, init_values=1.0, img_size=518,
     )
-
     model = _create_vision_transformer(
         'vit_base_patch14_dinov2', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
@@ -1932,14 +1998,13 @@ def vit_large_patch14_dinov2(pretrained=False, **kwargs):
     """ ViT-L/14 for DINOv2
     """
     model_args = dict(
-        patch_size=14, embed_dim=1024, depth=24, num_heads=16, 
-        init_values=1.0, img_size=518, 
+        patch_size=14, embed_dim=1024, depth=24, num_heads=16, init_values=1.0, img_size=518,
     )
-
     model = _create_vision_transformer(
         'vit_large_patch14_dinov2', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
+
 @register_model
 def vit_giant_patch14_dinov2(pretrained=False, **kwargs):
     """ ViT-G/14 for DINOv2
@@ -1952,13 +2017,13 @@ def vit_giant_patch14_dinov2(pretrained=False, **kwargs):
 
     model_args = dict(
         patch_size=14, embed_dim=1536, depth=40, num_heads=24, init_values=1.0, 
-        mlp_ratio=2.66667 * 2, ffn_layer=SwiGLUPacked, img_size=518, act_layer=nn.SiLU
+        mlp_ratio=2.66667 * 2, mlp_layer=SwiGLUPacked, img_size=518, act_layer=nn.SiLU
     )
-
     model = _create_vision_transformer(
         'vit_giant_patch14_dinov2', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
+
 register_model_deprecations(__name__, {
     'vit_tiny_patch16_224_in21k': 'vit_tiny_patch16_224.augreg_in21k',
     'vit_small_patch32_224_in21k': 'vit_small_patch32_224.augreg_in21k',
diff --git a/timm/models/vision_transformer_hybrid.py b/timm/models/vision_transformer_hybrid.py
index f41e90e89b..8cf7bec1e6 100644
--- a/timm/models/vision_transformer_hybrid.py
+++ b/timm/models/vision_transformer_hybrid.py
@@ -14,6 +14,7 @@
 Hacked together by / Copyright 2020, Ross Wightman
 """
 from functools import partial
+from typing import List, Tuple
 
 import torch
 import torch.nn as nn
@@ -74,10 +75,43 @@ def forward(self, x):
         x = self.backbone(x)
         if isinstance(x, (list, tuple)):
             x = x[-1]  # last feature if backbone outputs list/tuple of features
-        x = self.proj(x).flatten(2).transpose(1, 2)
+        x = self.proj(x)
+        x = x.flatten(2).transpose(1, 2)
         return x
 
 
+class HybridEmbedWithSize(nn.Module):
+    """ CNN Feature Map Embedding
+    Extract feature map from CNN, flatten, project to embedding dim.
+    """
+    def __init__(
+            self,
+            backbone,
+            img_size=224,
+            patch_size=1,
+            feature_size=None,
+            in_chans=3,
+            embed_dim=768,
+            bias=True,
+    ):
+        super().__init__(
+            backbone=backbone,
+            img_size=img_size,
+            patch_size=patch_size,
+            feature_size=feature_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+            bias=bias,
+        )
+
+    def forward(self, x) -> Tuple[torch.Tensor, List[int]]:
+        x = self.backbone(x)
+        if isinstance(x, (list, tuple)):
+            x = x[-1]  # last feature if backbone outputs list/tuple of features
+        x = self.proj(x)
+        return x.flatten(2).transpose(1, 2), x.shape[-2:]
+
+
 def _create_vision_transformer_hybrid(variant, backbone, pretrained=False, **kwargs):
     embed_layer = partial(HybridEmbed, backbone=backbone)
     kwargs.setdefault('patch_size', 1)  # default patch size for hybrid models if not set

From 960a8825105096625511853ab211a8e5ab35b96d Mon Sep 17 00:00:00 2001
From: Ross Wightman <rwightman@gmail.com>
Date: Tue, 9 May 2023 18:00:41 -0700
Subject: [PATCH 6/7] Remove label offsets and remove old weight url for 1001
 class (background + in1k) TF origin weights

---
 timm/models/inception_resnet_v2.py |  4 ----
 timm/models/inception_v3.py        | 11 ++---------
 timm/models/inception_v4.py        |  2 --
 timm/models/nasnet.py              |  1 -
 timm/models/pnasnet.py             |  2 --
 5 files changed, 2 insertions(+), 18 deletions(-)

diff --git a/timm/models/inception_resnet_v2.py b/timm/models/inception_resnet_v2.py
index 138470f750..f4efaf520d 100644
--- a/timm/models/inception_resnet_v2.py
+++ b/timm/models/inception_resnet_v2.py
@@ -314,23 +314,19 @@ def _create_inception_resnet_v2(variant, pretrained=False, **kwargs):
     # ported from http://download.tensorflow.org/models/inception_resnet_v2_2016_08_30.tar.gz
     'inception_resnet_v2.tf_in1k': {
         'hf_hub_id': 'timm/',
-        'url': 'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/inception_resnet_v2-940b1cd6.pth',
         'num_classes': 1000, 'input_size': (3, 299, 299), 'pool_size': (8, 8),
         'crop_pct': 0.8975, 'interpolation': 'bicubic',
         'mean': IMAGENET_INCEPTION_MEAN, 'std': IMAGENET_INCEPTION_STD,
         'first_conv': 'conv2d_1a.conv', 'classifier': 'classif',
-        'label_offset': 1,  # 1001 classes in pretrained weights
     },
     # As per https://arxiv.org/abs/1705.07204 and
     # ported from http://download.tensorflow.org/models/ens_adv_inception_resnet_v2_2017_08_18.tar.gz
     'inception_resnet_v2.tf_ens_adv_in1k': {
         'hf_hub_id': 'timm/',
-        'url': 'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/ens_adv_inception_resnet_v2-2592a550.pth',
         'num_classes': 1000, 'input_size': (3, 299, 299), 'pool_size': (8, 8),
         'crop_pct': 0.8975, 'interpolation': 'bicubic',
         'mean': IMAGENET_INCEPTION_MEAN, 'std': IMAGENET_INCEPTION_STD,
         'first_conv': 'conv2d_1a.conv', 'classifier': 'classif',
-        'label_offset': 1,  # 1001 classes in pretrained weights
     }
 })
 
diff --git a/timm/models/inception_v3.py b/timm/models/inception_v3.py
index 984970817d..1c2a407ec0 100644
--- a/timm/models/inception_v3.py
+++ b/timm/models/inception_v3.py
@@ -428,21 +428,14 @@ def _cfg(url='', **kwargs):
         hf_hub_id='timm/',
         url='https://download.pytorch.org/models/inception_v3_google-1a9a5a14.pth'),
     # my port of Tensorflow SLIM weights (http://download.tensorflow.org/models/inception_v3_2016_08_28.tar.gz)
-    'inception_v3.tf_in1k': _cfg(
-        hf_hub_id='timm/',
-        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_inception_v3-e0069de4.pth',
-        num_classes=1000, label_offset=1),
+    'inception_v3.tf_in1k': _cfg(hf_hub_id='timm/'),
     # my port of Tensorflow adversarially trained Inception V3 from
     # http://download.tensorflow.org/models/adv_inception_v3_2017_08_18.tar.gz
-    'inception_v3.tf_adv_in1k': _cfg(
-        hf_hub_id='timm/',
-        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/adv_inception_v3-9e27bd63.pth',
-        num_classes=1000, label_offset=1),
+    'inception_v3.tf_adv_in1k': _cfg(hf_hub_id='timm/'),
     # from gluon pretrained models, best performing in terms of accuracy/loss metrics
     # https://gluon-cv.mxnet.io/model_zoo/classification.html
     'inception_v3.gluon_in1k': _cfg(
         hf_hub_id='timm/',
-        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/gluon_inception_v3-9f746940.pth',
         mean=IMAGENET_DEFAULT_MEAN,  # also works well with inception defaults
         std=IMAGENET_DEFAULT_STD,  # also works well with inception defaults
     )
diff --git a/timm/models/inception_v4.py b/timm/models/inception_v4.py
index 91c435d978..a43290a3db 100644
--- a/timm/models/inception_v4.py
+++ b/timm/models/inception_v4.py
@@ -312,12 +312,10 @@ def _create_inception_v4(variant, pretrained=False, **kwargs) -> InceptionV4:
 default_cfgs = generate_default_cfgs({
     'inception_v4.tf_in1k': {
         'hf_hub_id': 'timm/',
-        'url': 'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-cadene/inceptionv4-8e4777a0.pth',
         'num_classes': 1000, 'input_size': (3, 299, 299), 'pool_size': (8, 8),
         'crop_pct': 0.875, 'interpolation': 'bicubic',
         'mean': IMAGENET_INCEPTION_MEAN, 'std': IMAGENET_INCEPTION_STD,
         'first_conv': 'features.0.conv', 'classifier': 'last_linear',
-        'label_offset': 1,  # 1001 classes in pretrained weights
     }
 })
 
diff --git a/timm/models/nasnet.py b/timm/models/nasnet.py
index edf355f8e0..954ee176b0 100644
--- a/timm/models/nasnet.py
+++ b/timm/models/nasnet.py
@@ -588,7 +588,6 @@ def _create_nasnet(variant, pretrained=False, **kwargs):
         'num_classes': 1000,
         'first_conv': 'conv0.conv',
         'classifier': 'last_linear',
-        'label_offset': 1,  # 1001 classes in pretrained weights
     },
 })
 
diff --git a/timm/models/pnasnet.py b/timm/models/pnasnet.py
index 23785007c6..bee18604c3 100644
--- a/timm/models/pnasnet.py
+++ b/timm/models/pnasnet.py
@@ -355,7 +355,6 @@ def _create_pnasnet(variant, pretrained=False, **kwargs):
 default_cfgs = generate_default_cfgs({
     'pnasnet5large.tf_in1k': {
         'hf_hub_id': 'timm/',
-        'url': 'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-cadene/pnasnet5large-bf079911.pth',
         'input_size': (3, 331, 331),
         'pool_size': (11, 11),
         'crop_pct': 0.911,
@@ -365,7 +364,6 @@ def _create_pnasnet(variant, pretrained=False, **kwargs):
         'num_classes': 1000,
         'first_conv': 'conv_0.conv',
         'classifier': 'last_linear',
-        'label_offset': 1,  # 1001 classes in pretrained weights
     },
 })
 

From b9d43c7dcac1fe05e851dd7be7187b108af593d2 Mon Sep 17 00:00:00 2001
From: Ross Wightman <rwightman@gmail.com>
Date: Tue, 9 May 2023 20:38:10 -0700
Subject: [PATCH 7/7] Version 0.8.22dev0

---
 timm/version.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/timm/version.py b/timm/version.py
index b7b0e31206..915315d246 100644
--- a/timm/version.py
+++ b/timm/version.py
@@ -1 +1 @@
-__version__ = '0.8.21dev0'
+__version__ = '0.8.22dev0'