GitLab

Switch line orders in trainer so that restore_map is called after moving average variables are created.

Moving averages are now properly loaded during fine-tuning, instead of being recreated.

PiperOrigin-RevId: 190496046