import itertools


def train_iterator(num_iterations):
    """Iterator that returns an index of the current step.
  This iterator runs forever if num_iterations is None
  otherwise it runs for some fixed amount of steps.
  """
    if num_iterations is None:
        return itertools.count()
    else:
        return range(num_iterations)