def step_lr(learning_rate, epoch):
    """Step Learning rate.

  Args:
    learning_rate: Initial learning rate.
    epoch: Current epoch we are one. This is one based.

  Returns:
    The learning rate to be used for this current batch.
  """
    if epoch < 80:
        return learning_rate
    elif epoch < 120:
        return learning_rate * 0.1
    else:
        return learning_rate * 0.01