How do I load a pre-trained model?

[btaba]

Discussed in https://github.com/google/brax/discussions/403

^{Originally posted by eleninisioti October 11, 2023} There is a notebook that explains how to save and load models (https://github.com/google/brax/blob/main/notebooks/training.ipynb) but there testing happens right after training, calling function make_inference_fn(params), which requires first running make_inference_fn, params, _ = train_fn(environment=env, progress_fn=progress), the whole training process.

My question is how I can test without running the training process, simply by loading the params. How can I have inference_fn withour training?