mega vit

kyegomez · Oct 3, 2023 · c1f6f63 · c1f6f63
1 parent 5e8e1ed
commit c1f6f63
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/mega_vit/__init__.py b/mega_vit/__init__.py
@@ -1 +1 @@
-from mega_vit.main import MegaVit
+from mega_vit.main import MegaVit   
diff --git a/mega_vit/main.py b/mega_vit/main.py
@@ -10,6 +10,16 @@ def pair(t):
 
 # classes
 
+class Parallel(nn.Module):
+    def __init__(self, *fns):
+        super().__init__()
+        self.fns = nn.ModuleList(fns)
+
+    def forward(self, x):
+        return ([fn(x) for fn in self.fns])
+
+
+
 class LayerNorm(nn.Module):
     def __init__(
         self,
@@ -104,7 +114,8 @@ def forward(self, x):
 
             #project out
             return self.to_out(out)
-
+
+
 class Transformer(nn.Module):
     def __init__(
         self,