e9b32747b4
* 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 |
||
|---|---|---|
| .. | ||
| main.py | ||
| mnist.py | ||
| README.md | ||
| requirements.txt | ||
MNIST
This example shows how to run some simple models on MNIST.
Install the dependencies:
pip install -r requirements.txt
Run the example with:
python main.py
By default, the example runs on the CPU. To run on the GPU, use:
python main.py --gpu
For a full list of options run:
python main.py --help
To run the PyTorch or JAX examples install the respective framework.