Given an input, I would like to do multiple classification tasks.
Each input needs to be classified into one of 5 classes. There are 6 such classification tasks to be done. I could build six separate Linear(some_number, 5) layers and return the result as tuple in the forward() function. Then call the loss function 6 times and sum the losses to produce the overall loss.
I am wondering if I could do this better than this. For example, can I have a single Linear(some_number, 5*6) as the output. Then reshape the logits to (6,5) and use. The documentation for CrossEntropyLoss mentions about “K-dimensional loss”. Is it for use cases as here? If so, can anyone share a sample code (When I tried to use I got confused about how to reshape the output).
https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html 2