Track more than one loss using Trainer and Wandb

Intermediate
1

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(
    loss=loss,
    ...)

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