zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-06-24 01:17:28 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity
main
mlx-examples/transformer_lm/README.md
Angelos Katharopoulos e9b32747b4
Add grad checkpointing and PE in the transformer example (#387) 
* Add grad checkpointing and PE in the transformer example

* Remove other frameworks from LM example

* Remove the other frameworks from MNIST example

* Improve the transformer LM example

* Fix black and change LR
2024-02-01 13:04:03 -08:00

316 B
Raw Permalink Blame History

Transformer LM

This is an example of a decoder-only Transformer LM. The only dependency is MLX.

Run the example on the GPU with:

python main.py --gpu

By default the dataset is the PTB corpus. Choose a different dataset with the --dataset option.