def chunk_tensors()

in lm_human_preferences/utils/core.py [0:0]

• 11 lines of code
• 4 McCabe index (conditional complexity)

def chunk_tensors(tensors, *, limit=1 << 28):
    """Chunk the list of tensors into groups of size at most `limit` bytes.

    The tensors must have a static shape.
    """
    total = 0
    batches = []
    for v in tensors:
        size = v.dtype.size * v.shape.num_elements()
        if not batches or total + size > limit:
            total = 0
            batches.append([])
        total += size
        batches[-1].append(v)
    return batches