clean up

kyegomez · Sep 27, 2023 · 594e720 · 594e720
1 parent ec5f2ad
commit 594e720
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/mmca/main.py b/mmca/main.py
@@ -1,7 +1,8 @@
-import torch 
-from torch import nn
-from einops import rearrange
+import torch
 import torch.nn.functional as F
+from einops import rearrange
+from torch import nn
+
 # from zeta.nn import FlashAttention
 
 class MultiModalCausalAttention(nn.Module):
@@ -103,6 +104,7 @@ def __init__(
             embed_dim=dim, 
             num_heads=heads
         )
+
         self.cross_attn = nn.MultiheadAttention(
             embed_dim=dim, 
             num_heads=heads

diff --git a/simple_example.py b/simple_example.py
@@ -1,7 +1,6 @@
 import torch 
 from mmca.main import SimpleMMCA
 
-
 # Define the dimensions
 dim = 512
 head = 8