Track more than one loss using Trainer and Wandb

I’ve been searching for this information in the forum for a while now and haven’t found any satisfactory answers. I’m currently training a model with the Hugging Face Trainer and tracking the loss with Wandb. In the forward function of my model, I have something like this:

loss = loss_1 + loss_2
return ModelOutput(

What would be the best way to track all the losses here? Currently, Wandb only tracks the one included in the model output.

Try to use my implementation: It is easy to use.