Let’s build a NN with JAX

Load dataset

Among the many options to load a dataset, let’s use Hugging Face:

Define model architecture

Here, we use Flax:

Train

Checkpointing

Test

Save model