Fix initialization of training statistics

This fixes a bug introduced in commit ac6fc8bf, where training
statistics are initialized *after* loading values from a checkpoint,
rather than before. Due to this bug, parts of the checkpoint state are
effectively lost.