def METHOD_NAME(di, args):
    current_epoch = -1
    logger.info('{}'.format(di.size))
    pbar = None
    for data in di:
        time.sleep(args.wait / 1000.0)
        if di.epoch >= args.max_epoch:
            break
        if current_epoch != di.epoch:
            current_epoch = di.epoch
            if pbar is not None:
                pbar.close()
            logger.info('Epoch {}'.format(current_epoch))
            pbar = tqdm(total=di.size)
        pbar.update(len(data[0]))