zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-08-22 13:07:51 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

run precommit

Browse Source
This commit is contained in:
junwoo-yun 2023-12-25 07:52:19 +08:00
parent 0e0557b756
commit 88a94b9db8
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -156,13 +156,13 @@ class TransformerEncoder(Module):
num_heads: int,
mlp_dims: Optional[int] = None,
dropout: float = 0.0,
norm_first: bool = False,
activation=relu,
norm_first: bool = False,
):
super().__init__()
self.layers = [
TransformerEncoderLayer(
dims, num_heads, mlp_dims, dropout, norm_first, activation
dims, num_heads, mlp_dims, dropout, activation, norm_first
)
for i in range(num_layers)
]
@ -246,13 +246,13 @@ class TransformerDecoder(Module):
num_heads: int,
mlp_dims: Optional[int] = None,
dropout: float = 0.0,
norm_first: bool = False,
activation=relu,
norm_first: bool = False,
):
super().__init__()
self.layers = [
TransformerDecoderLayer(
dims, num_heads, mlp_dims, dropout, norm_first, activation
dims, num_heads, mlp_dims, dropout, activation, norm_first
)
for i in range(num_layers)
]