Target size (torch.Size([8])) must be the same as input size (torch.Size([8, 2]))

What if the labels are already typed as int64
I have the same error for sequence classification with BERT model, even my labels start from 0, here’s what my training dataset looks like

Dataset({
    features: ['label', 'text', 'input_ids', 'token_type_ids', 'attention_mask'],
    num_rows: 393
})