def multicrop_collate_fn(samples):
    """Multi-crop collate function for VISSL integration.

    Run custom collate on a single key since VISSL transforms affect only DataKeys.INPUT

    """
    result = vissl_collate_helper(samples)

    inputs = [[] for _ in range(len(samples[0][DataKeys.INPUT]))]
    for batch_ele in samples:
        multi_crop_imgs = batch_ele[DataKeys.INPUT]

        for idx, crop in enumerate(multi_crop_imgs):
            inputs[idx].append(crop)

    for idx, ele in enumerate(inputs):
        inputs[idx] = torch.stack(ele)

    result[DataKeys.INPUT] = inputs

    return result
Multi-crop collate function for VISSL integration.

    Run custom collate on a single key since VISSL transforms affect only DataKeys.INPUT

    
src/flash/image/embedding/vissl/transforms/utilities.py