Hydra support

[bmacadam-sfu]

Thanks for putting this out there! I've been complaining about the lack of a Lightning-style trainer for JAX for years now, and was in the process of adapting an earlier version of this project from the UvA tutorials for my purposes. It seems like you're handling configs yourself here, and I'd like to suggest offloading that to hydra (you can look at this template to see how nice hydra can be in practice).

This is just a suggestion, of course, but I think it would make it easier for people to port their torch projects to jax. I'm also available to help out if you're interested.