Return a dummy data loader which returns a fixed data batch of target shape

def get_dummy_dataloader(dataloader, target_shape):
    """Return a dummy data loader which returns a fixed data batch of target shape"""
    data_iter = enumerate(dataloader)
    _, data_batch = next(data_iter)
    logging.debug('Searching target batch shape: %s', target_shape)
    while data_batch[0].shape != target_shape:
        logging.debug('Skip batch with shape %s', data_batch[0].shape)
        _, data_batch = next(data_iter)
    logging.debug('Found target dummy batch.')

    class DummyIter():
        def __init__(self, batch):
            self._batch = batch

        def __iter__(self):
            while True:
                yield self._batch

    return DummyIter(data_batch)