def simclr_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 = []
    num_views = len(samples[0][DataKeys.INPUT])
    view_idx = 0
    while view_idx < num_views:
        for batch_ele in samples:
            imgs = batch_ele[DataKeys.INPUT]
            inputs.append(imgs[view_idx])

        view_idx += 1

    result[DataKeys.INPUT] = torch.stack(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