this line:
weighted_loss = (unweighted_loss * weights).mean(dim=1)
should be corrected to:
weighted_loss = (unweighted_loss * weights).sum(dim=1)

reason:
when use mean, the padding locations will be calculated as denoimator to drag down the loss. In this case, model will learn to cheat by predicting as long as possible (as a result, no eos will be generated). ( I've tested it).
Also, use mean is inconsistent with the downstream loss calculation. In the downstream, loss will be divided by num_token. In this case, CEloss will be in total divided by square of num_token, which makes no sense.

when use sum, the padding locations will not be calculated (all zeros). In total loss will not be divieded by num_token square but just num_token.

thanks.