Disclaimer: This article is aimed at researchers and people who would like to publish their project in a no-nonsense way and definitely not aimed at seasoned developers who’re struggling to cope between stackoverflow and some coding agent. Do what you will with this information.

Pre-trained models are great. They’re trained on a lot of data us normies probably won’t be able to compile by ourselves and they also require a lot of compute to train from scratch. Ever since BERT was released, the NLP community has been using pre-trained models to fine-tune on their own datasets. This is a great way to leverage the power of these models without having to train them from scratch.

So say you’ve trained a model using flax, it trained fine, has a nice learning curve (train vs validation) and now you want to save it. Or, you want to save checkpoints of the model during specific stages of the training process and later, use the best checkpoints for inference. Technically, all flax modules are dataclasses and params (part of model state in flax) are what store the model, so what we need to do for checkpointing is to persist the params.

Update on October 2026: I have automated the post-installation with ansible. You can check out the ansible-playbooks repository.

Okay why not docker though?