RuntimeError: Expected tensor [2, 512, 1], src [2, 371, 6] and index [2, 512, 1] to have the same size apart from dimension 2

I use my own dataseta, and in the training process of CRF, the following error occurs:

```
line 483, in main
global_step, tr_loss = train(args, train_dataset, model, tokenizer, labels,pos, pad_token_label_id)

line 138, in train
loss = model.score(batch) # model outputs are always tuple in pytorch-transformers (see doc)

line 52, in score
return self.crf.score(output, labels_mask, labels)

line 46, in score
gold_score = self.crf.calc_gold_score(logits, labels, lens)

line 99, in calc_gold_score
unary_score = self.calc_unary_score(logits, labels, lens).sum(

line 93, in calc_unary_score
scores = torch.gather(logits, 2, labels_exp).squeeze(-1)

RuntimeError: Expected tensor [2, 512, 1], src [2, 371, 6] and index [2, 512, 1] to have the same size apart from dimension 2
```