update vit

marsggbo · Dec 12, 2022 · c4cdf95 · c4cdf95
1 parent b2f033a
commit c4cdf95
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/hyperbox/networks/vit/__init__.py b/hyperbox/networks/vit/__init__.py
@@ -0,0 +1 @@
+from .vit import *
diff --git a/hyperbox/networks/vit/vit.py b/hyperbox/networks/vit/vit.py
@@ -13,6 +13,19 @@
 from hyperbox.mutables import spaces, ops
 
 
+__all__ = [
+    'Attention',
+    'Transformer',
+    'VisionTransformer',
+    'ViT',
+    'ViT_S',
+    'ViT_B',
+    'ViT_H',
+    'ViT_G',
+    'ViT_10B',
+]
+
+
 # helpers
 def pair(t):
     return t if isinstance(t, tuple) else (t, t)