Fix race in train.py for PMap model, following checkpoint_manager update.
Fix race in train.py for PMap model, following checkpoint_manager update.

PiperOrigin-RevId: 418971983

tensorflow/lingvo,tensorflow/lingvo,tensorflow/lingvo,tensorflow/lingvo