model save to and load from a directionary?

[DonghanHe]

Hey I dunno if I am searching wrong or what, but how do you save flax's nn.Module models? I tried jnp.save('...', frozen_dict_of_auto_encoder_params) and the loaded stuff are just ['encoder', 'decoder'] and nothing else. Is there something like a tf.saved_models.save('...')?

[lamflokas]

This link shows an example on how to save and load FrozenDicts in Flax. The flax serialiation package contains a lot of helpful functions.

[DonghanHe]

awesome, thank you