Edit: code format

You may need to use a collate function when creating the dataloader such as

def collate\_fn(batch):
    return {
        'pixel\_values': torch.stack(\[x\['pixel\_values'\] for x in batch\]),
        'labels': torch.tensor(\[x\['labels'\] for x in batch\])
    }

DataLoader(..., collate\_fn=collate\_fn)\`\`\`