zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-08-29 16:58:45 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

add olmo2

Browse Source
This commit is contained in:
Awni Hannun 2024-11-27 12:15:45 -08:00
parent b0ef3a4c9e
commit fde4b4dc42
2 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
llms/mlx_lm/models/olmo2.py
View File

@ -231,7 +231,9 @@ class TransformerBlock(nn.Module):
self.post_attention_layernorm = nn.RMSNorm( self.post_attention_layernorm = nn.RMSNorm(
args.hidden_size, eps=args.rms_norm_eps args.hidden_size, eps=args.rms_norm_eps
) )
self.post_feedforward_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) self.post_feedforward_layernorm = nn.RMSNorm(
args.hidden_size, eps=args.rms_norm_eps
)
self.args = args self.args = args
def __call__( def __call__(

1
llms/tests/test_models.py
View File

@ -813,6 +813,5 @@ class TestModels(unittest.TestCase):
) )
if __name__ == "__main__": if __name__ == "__main__":
unittest.main() unittest.main()