ForwardKLWithChunkedOutputLoss should be averaged over total_non_pad_tokens

[moussaKam]

Hi there,

Currently the loss ForwardKLWithChunkedOutputLoss is averaged over the chunks which I think is not the right way to do it since there are chunks that contains only pad tokens. line

I think the right way to do it is to average over the total non pad tokens. I did the fix here but couldn't open a pull request because my commit didn't pass the test since the loss did not match with the expected loss.

[pbontrager]

Thank you for pointing this out. You should be able to submit a PR even if the tests don't pass. Before landing the PR, you'll need to verify that training is improved and then you can update the loss values in the test too.