def _set_train_steps(max_steps, save_checkpoints_steps, estimator):
    """ Set the train steps for each iteration in the train/eval loop. If the estimator
        has been trained before, the initial step will be updated accordingly.

    Args:
        max_steps: (int) number of steps to run training.
        save_checkpoints_steps: (int) frequency in steps to save checkpoints.
        estimator: tf.Estimator.estimator initialized object.

    Returns:
        list containing the train steps for each iteration of the train/eval loop.
    """

    # Get initial steps from *estimator* if a model has already been saved.
    try:
        initial_step = estimator.get_variable_value("global_step")
        remain_steps = max_steps - initial_step
        num_loops, remain = divmod(remain_steps, save_checkpoints_steps)

    except ValueError:
        initial_step = 0
        num_loops, remain = divmod(max_steps, save_checkpoints_steps)

    train_steps = [initial_step + (i * save_checkpoints_steps) for i in range(1, num_loops + 1)]
    if remain > 0:
        train_steps.append(train_steps[-1] + remain)

    return train_steps