updated train for computing per token loss

[adamimos]

had to explicitly calculate loss for this

[melembroucarlitos]

I'm confused. Why are we calculating loss per token if we're not logging the value?

[adamimos]

The main thing I was thinking is that this was the more general thing, that would work for cases where we did or did not want to log by token. I was going to implement the logging by token in a different pull request.

A secondary issue was that the way transformerlens deals with loss is to not compute it on the final token since there's no data for what that prediction should be. So I at least wanted to manually do cross-entropy loss for the logits in order to get that (but this is a seperate issue for the per token thing).