first, i don’t know if the fallownig part is right.



 Lady_Hangaku:

 # (batchsize, num of tokens, num of classes)


base on :

it should be (batchsize, num of classes, num of tokens).
second ,pred_output must be log of probability vector.
https://pytorch.org/docs/stable/generated/torch.nn.LogSoftmax.html#torch.nn.LogSoftmax