fix: Fix wrong layer norm used in TransformerDecoderLayer (#171)

* fix wrong ln used in transformer decoder layer * run pre-commit
ml-explore · Dec 4, 2024 · d5fefbf · d5fefbf
1 parent 30eecff
commit d5fefbf
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/Source/MLXNN/Transformer.swift b/Source/MLXNN/Transformer.swift
@@ -267,7 +267,7 @@ class TransformerDecoderLayer: Module {
 
             y = crossAttention(y, keys: memory, values: memory, mask: memoryMask)
             y = dropout2(y)
-            x = ln1(x + y)
+            x = ln2(x + y)
 
             y = linear1(x)
             y = activation(y)