def _get_input_iterator(input_fn, strategy):
    """Returns distributed dataset iterator."""
    input_data = input_fn()
    if callable(input_data):
        iterator = iter(strategy.
            experimental_distribute_datasets_from_function(input_data))
    else:
        iterator = iter(strategy.experimental_distribute_dataset(input_data))
    return iterator
