import logging


def to_generator(train_step_fn):
    """
        Makes the provided train step function a Python generator.
        The stop condition for the returned generator is when the
        train_step_fn raises a ``StopIteration``, which is interpreted
        as end of data.
    """

    def wrap(state):
        while True:
            try:
                yield train_step_fn(state)
            except StopIteration as e:
                logging.info("End of data reached from train_step", exc_info=e)
                return

    return wrap