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

[possible1402]

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