A custom trainer for multi-task learning?

I am working on using Trainer under a multi-task setting. Specifically, a list of losses([loss1, loss2, ...]) is returned in a single model forward, and optimized with a custom optimizer like PCGrad. How can I do this with minimal changes to Trainer (while preserving all the nice features of Trainer like multi-gpu training)?

1 Like