Question about Inference

[SeeYaaa]

Is there a way to save the model and use it in inference?

[SeeYaaa]

Is it possible to use orbax hier? Like this: https://github.com/google/orbax/blob/main/export/orbax/export/examples/periodically_export_main.py