From 081e6c2b9642d74cc8b85d18f0ea9272f6af65ec Mon Sep 17 00:00:00 2001
From: Ryan <23580140+brianhou0208@users.noreply.github.com>
Date: Tue, 29 Apr 2025 05:19:53 +0800
Subject: [PATCH] update swiftformer

---
 timm/models/swiftformer.py | 34 +++++++++++++++++-----------------
 1 file changed, 17 insertions(+), 17 deletions(-)

diff --git a/timm/models/swiftformer.py b/timm/models/swiftformer.py
index a349a2e7..748569b4 100644
--- a/timm/models/swiftformer.py
+++ b/timm/models/swiftformer.py
@@ -480,7 +480,7 @@ class SwiftFormer(nn.Module):
             indices: Union[int, List[int]] = 1,
             prune_norm: bool = False,
             prune_head: bool = True,
-    ) -> List[int]:
+    ):
         """ Prune layers not required for specified intermediates.
         """
         take_indices, max_index = feature_take_indices(len(self.stages), indices)
@@ -561,10 +561,10 @@ def _cfg(url: str = '', **kwargs: Any) -> Dict[str, Any]:
 
 
 default_cfgs = generate_default_cfgs({
-    # 'swiftformer_xs.dist_in1k': _cfg(url=''),
-    # 'swiftformer_s.dist_in1k': _cfg(url=''),
-    # 'swiftformer_l1.dist_in1k': _cfg(url=''),
-    # 'swiftformer_l3.dist_in1k': _cfg(url=''),
+    # 'swiftformer_xs.dist_in1k': _cfg(hf_hub_id='timm/'),
+    # 'swiftformer_s.dist_in1k': _cfg(hf_hub_id='timm/'),
+    # 'swiftformer_l1.dist_in1k': _cfg(hf_hub_id='timm/'),
+    # 'swiftformer_l3.dist_in1k': _cfg(hf_hub_id='timm/'),
     'swiftformer_xs.untrained': _cfg(),
     'swiftformer_s.untrained': _cfg(),
     'swiftformer_l1.untrained': _cfg(),
@@ -588,18 +588,18 @@ def swiftformer_xs(pretrained: bool = False, **kwargs: Any) -> SwiftFormer:
     return _create_swiftformer('swiftformer_xs', pretrained=pretrained, **dict(model_args, **kwargs))
 
 
-# @register_model
-# def swiftformer_s(pretrained: bool = False, **kwargs: Any) -> SwiftFormer:
-#     model_args = dict(layers=[3, 3, 9, 6], embed_dims=[48, 64, 168, 224])
-#     return _create_swiftformer('swiftformer_s', pretrained=pretrained, **dict(model_args, **kwargs))
+@register_model
+def swiftformer_s(pretrained: bool = False, **kwargs: Any) -> SwiftFormer:
+    model_args = dict(layers=[3, 3, 9, 6], embed_dims=[48, 64, 168, 224])
+    return _create_swiftformer('swiftformer_s', pretrained=pretrained, **dict(model_args, **kwargs))
 
-# @register_model
-# def swiftformer_l1(pretrained: bool = False, **kwargs: Any) -> SwiftFormer:
-#     model_args = dict(layers=[4, 3, 10, 5], embed_dims=[48, 96, 192, 384])
-#     return _create_swiftformer('swiftformer_l1', pretrained=pretrained, **dict(model_args, **kwargs))
+@register_model
+def swiftformer_l1(pretrained: bool = False, **kwargs: Any) -> SwiftFormer:
+    model_args = dict(layers=[4, 3, 10, 5], embed_dims=[48, 96, 192, 384])
+    return _create_swiftformer('swiftformer_l1', pretrained=pretrained, **dict(model_args, **kwargs))
 
 
-# @register_model
-# def swiftformer_l3(pretrained: bool = False, **kwargs: Any) -> SwiftFormer:
-#     model_args = dict(layers=[4, 4, 12, 6], embed_dims=[64, 128, 320, 512])
-#     return _create_swiftformer('swiftformer_l3', pretrained=pretrained, **dict(model_args, **kwargs))
\ No newline at end of file
+@register_model
+def swiftformer_l3(pretrained: bool = False, **kwargs: Any) -> SwiftFormer:
+    model_args = dict(layers=[4, 4, 12, 6], embed_dims=[64, 128, 320, 512])
+    return _create_swiftformer('swiftformer_l3', pretrained=pretrained, **dict(model_args, **kwargs))
\ No newline at end of file