New version of flax requires frozen dicts
Small update to maintain compatibility with new version of flax.

@mfuntowicz 
Hey @KristianHolsheimer,

Sorry I saw your PR a bit too late...I think this is solved now on master no?<|||||>Okay no worries. Sorry I should've tagged you. Thanks for the reply