Hi all,
I am faced with the following situation. I am using one model to solve multiple classification tasks, where each classification task itself is multi-class, and the number of possible classes varies across classification tasks. To give an example:
The model outputs a vector with 22 elements, where I would like to apply a softmax over:

The first 5 elements
The following 5 elements
The following 8 elements
The last 4 elements

This is because the model is simultaneously solving 4 classification tasks, where the first 2 tasks have 5 candidate classes each, the third task has 8 candidate classes, and the final task has 4 candidate classes.
I would also like to define an appropriate cross-entropy loss that follows this same structure.
My questions are:

How can I use torch.nn.Softmax to achieve this?
How can I define the custom cross-entropy loss mentioned above?

Many thanks!