Make post_process_all_at_once schedule param possibly dataset-specific
Make post_process_all_at_once schedule param possibly dataset-specific

PiperOrigin-RevId: 407325234

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