Multiple Classification Heads (For two tier labelling)

Hello,

I have a dataset that has two levels of labels. They are coarse and fine labels. The original paper mentions that it uses multiple classification heads so that it can train on both the coarse and fine labels.

As I have limited compute, I rely on the optimisations made available through the HF Trainer. However, when I try to pass two custom named label columns, I run into some issues. I’m thinking there must be a simpler approach to this. Could anyone advise me on adding multiple classification heads to my model? And in particular how to do this using the Trainer class.