Use orbax for saving trained models

[grahamgower]

https://github.com/google/orbax

Maybe this allows use to remove some code? Also related to #169.