Merge remote-tracking branch 'origin/main' into master-patched

xwang233 · Aug 7, 2024 · 3a8ddc5 · 3a8ddc5
2 parents f354fc9 + e9ef942
commit 3a8ddc5
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 0 deletions.
diff --git a/timm/models/efficientformer_v2.py b/timm/models/efficientformer_v2.py
@@ -29,6 +29,8 @@
 from ._registry import generate_default_cfgs, register_model
 
 
+__all__ = ['EfficientFormerV2']
+
 EfficientFormer_width = {
     'L': (40, 80, 192, 384),  # 26m 83.3% 6attn
     'S2': (32, 64, 144, 288),  # 12m 81.6% 4attn dp0.02

diff --git a/timm/models/fastvit.py b/timm/models/fastvit.py
@@ -20,6 +20,7 @@
 from ._manipulate import checkpoint_seq
 from ._registry import register_model, generate_default_cfgs
 
+__all__ = ['FastVit']
 
 def num_groups(group_size, channels):
     if not group_size:  # 0 or None

diff --git a/timm/models/hiera.py b/timm/models/hiera.py
@@ -42,6 +42,9 @@
 from ._features_fx import register_notrace_function
 
 
+__all__ = ['Hiera']
+
+
 def conv_nd(n: int) -> Type[nn.Module]:
     """
     Returns a conv with nd (e.g., Conv2d for n=2). Work up to n=3.

diff --git a/timm/models/nextvit.py b/timm/models/nextvit.py
@@ -20,6 +20,8 @@
 from ._manipulate import checkpoint_seq
 from ._registry import generate_default_cfgs, register_model
 
+__all__ = ['NextViT']
+
 
 def merge_pre_bn(module, pre_bn_1, pre_bn_2=None):
     """ Merge pre BN to reduce inference runtime.