def moco_collate_fn(samples):
    'MOCO collate function for VISSL integration.\n\n    Run custom collate on a single key since VISSL transforms affect only DefaultDataKeys.INPUT\n    '
    result = vissl_collate_helper(samples)
    inputs = []
    for batch_ele in samples:
        inputs.append(torch.stack(batch_ele[DefaultDataKeys.INPUT]))
    result[DefaultDataKeys.INPUT] = torch.stack(inputs).squeeze()[:, 0, :, :, :].squeeze()
    result['data_momentum'] = torch.stack(inputs).squeeze()[:, 1, :, :, :].squeeze()
    return result
MOCO collate function for VISSL integration.

Run custom collate on a single key since VISSL transforms affect only DefaultDataKeys.INPUT
def moco_collate_fn(samples):
    'MOCO collate function for VISSL integration.\n\n    Run custom collate on a single key since VISSL transforms affect only DefaultDataKeys.INPUT\n    '
    result = vissl_collate_helper(samples)
    inputs = []
    for batch_ele in samples:
        inputs.append(torch.stack(batch_ele[DefaultDataKeys.INPUT]))
    result[DefaultDataKeys.INPUT] = torch.stack(inputs).squeeze()[:, 0, :, :, :].squeeze()
    result['data_momentum'] = torch.stack(inputs).squeeze()[:, 1, :, :, :].squeeze()
    return result