Fix calculation of logits for training loss

[bob80333]

The embeddings must be l2 normalized before calculating the logits. This fix corrects both the loss that is logged and the training loss. The validation loss was already correct, as the model forward function normalizes the embeddings before computing the logits.

[Zasder3]

Really good catch! My only change was using torch.nn.functional.normalize to shrink the code a little bit.