I meet a problem that the Pytorch Loss functions such as functional.nll_loss and nn.CrossEntropyLoss are not supporting subset dataset.

I want to splitted the mnist and cifar10 dataset for specific use. For example, while loading the mnist data, I explicitly exclude the data with lable β6β in my data and label set so that the training/test dataset only have number 0-5 and 7-9.

However, the loss function will fail
loss=F.nll_loss(output, target)
or
loss=nn.CrossEntropyLoss(output, target)
since the function expect:
input(N, C) where C = number of classes
Target: (N) where each value is 0 <= targets[i] <= C-1 http://pytorch.org/docs/master/nn.html

In my case, this will fail since my C=9 but my target is 0-5+7-9 in stead of 0-8.

It should be a tensor of size (N, C), where the element output[n, c] should be the probability that element n is of class c. Is this what it is? (It doesnβt sound like it, because output[n, 6] should be equal to 0 but the error is telling you you donβt have enough enough classes).

My output is in the right format, (N, C) = (64, 9). The issue is that since I have 9 classes, the loss function by default take/expects labels 0-8. However, since I only get rid of number β6β, the labels left is β1,2,3,4,5,7,8,9β. The β9β exceed the loss expectation.

An alternative would be to keep the labels as they are, to allow the model to predict any label from 0 to 9 and hope that the training process is sufficient to teach it never to predict 6.