softmax after kernel

kyegomez · Oct 3, 2023 · 5e8e1ed · 5e8e1ed
1 parent 8f9e005
commit 5e8e1ed
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/mega_vit/main.py b/mega_vit/main.py
@@ -92,6 +92,9 @@ def forward(self, x):
         with torch.backends.cuda.sdp_kernel(enable_math=True):
             #attention
             out = F.scaled_dot_product_attention(q, k, v)
+
+            #softmax
+            out = self.attend(out)
 
             #dropout
             out = self.dropout(out)