Removed the creation of trained params and eval keys from the eval setup method as it's not really the responsibility of that method, it should just make the evaluation methods.
Also created the unreplicate_batch_dim method to easily unreplicate our update_batch_size dim.
Finally removed the unecessary vmap inside the evaluator as flax will automatically vmap for us
What?
Minor clean up to the evaluator.
Removed the creation of trained params and eval keys from the eval setup method as it's not really the responsibility of that method, it should just make the evaluation methods.
Also created the
unreplicate_batch_dimmethod to easily unreplicate ourupdate_batch_sizedim.Finally removed the unecessary vmap inside the evaluator as flax will automatically vmap for us