import tensorflow as tf

            'end_learning_rate': FLAGS.end_learning_rate,
            'warmup_learning_rate': FLAGS.warmup_learning_rate,
            'warmup_steps': FLAGS.warmup_steps,
            'decay_boundaries': parse_comma_list(decay_boundaries),
            'lr_decay_factors': parse_comma_list(lr_decay_factors),
        })

    tf.gfile.MakeDirs(model_dir)
    tf.logging.info('Starting to train model {}.'.format(model_scope))
    for _ in range(train_epochs // epochs_per_eval):
        tensors_to_log = {
            'lr': 'learning_rate',
            'loss': 'total_loss',
            'mse': 'mse_loss',
