zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-24 17:31:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix transformer (#1327)

Browse Source
This commit is contained in:
Awni Hannun 2024-08-13 16:04:26 -07:00 committed by GitHub
parent eaaea02010
commit 63ae767232
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
python/mlx/nn/layers/__init__.py
View File

@ -72,6 +72,8 @@ from mlx.nn.layers.recurrent import GRU, LSTM, RNN
from mlx.nn.layers.transformer import ( from mlx.nn.layers.transformer import (
MultiHeadAttention, MultiHeadAttention,
Transformer, Transformer,
TransformerDecoder,
TransformerDecoderLayer,
TransformerEncoder, TransformerEncoder,
TransformerEncoderLayer, TransformerEncoderLayer,
) )

4
python/mlx/nn/layers/transformer.py
View File

@ -147,9 +147,9 @@ class TransformerEncoderLayer(Module):
else: else:
y = self.attention(x, x, x, mask) y = self.attention(x, x, x, mask)
y = self.dropout1(y) y = self.dropout1(y)
y = self.ln1(x + y) x = self.ln1(x + y)
y = self.linear1(y) y = self.linear1(x)
y = self.activation(y) y = self.activation(y)
y = self.dropout2(y) y = self.dropout2(y)
y = self.linear2(y) y = self.linear2(y)