Hi marton,
As you can see in the doc 23

You should apply argmax rather than softmax on the output
638 is the actual sequence length, usually we pad sentences just to the length of the longest sequence in the batch.

What you have got is called prediction_scores, now you can do
predicts = prediction_scores.argmax(dim=-1) # (batch_size, sequence_length) dtype=torch.Long
scores = prediction_scores[predicts] # (batch_size, sequence_length) dtype=torch.float

scores now is the probability of all predicted tokens, if you just want the first token of every sentence, just modify the code above.