def select_buffer(buffer, batch_slot, count):
    """Util function to select the relevent subpart of a buffer for a given
    batch_slot and batch size"""
    if buffer is None:
        return None
    if isinstance(buffer, tuple):
        return tuple(select_buffer(x, batch_slot, count) for x in buffer)

    return buffer[batch_slot][:count]
Util function to select the relevent subpart of a buffer for a given
    batch_slot and batch size
ffcv/loader/epoch_iterator.py