How to properly add an additional scene classification head.

What's the feature?

Not a feature request but a burning question.

I'm new to the library and struggling to figure out how to do something quite simple: add a scene classification head. That is, beyond the object classes and their bounding boxes, I'd like to classify the scene of the image (a categorical variable with 8 values).

I've amended the dataset to load "scene_label" for a batch. I've subclassed an existing head and it now outputs three values instead of two: cls_score, bbox_pred, scene_probs.

However, It isn't clear to me what to change in order to handle the new scene_probs, the scene_label and compute the scene loss. What other functions or classes need to be created or changed for all this to work?

None of the tutorials/docs make this clear. At least for me. ChatGPT was also no help.

Any advice would be appreciated.

Any other context?

No response