Upload DeCodon

modeling_decodon.py

@@ -629,7 +629,7 @@ class MultiHeadedSelfAttention(nn.Module):
             attention_probs = attention_scores.softmax(dim=-1)
             attention_probs = self.dropout(attention_probs)
 
-            x = einsum(attention_probs, v, "b h q k, b v h d -> b q h d")
+            x = einsum(attention_probs, v, "b h q k, b k h d -> b q h d")
 
         x = rearrange(x, "b q h d -> b q (h d)", h=self.num_heads)