I guess you might be filtering out N classes, which don’t have the class indices [0, N-1], but might have larger values, which would create errors while trying to calculate the loss in e.g. nn.CrossEntropyLoss.
This loss function expects the model output in the shape [batch_size, nb_classes, *] and a target tensor in the shape [batch_size, *] containing class indices in [0, nb_classes-1].
If you are e.g. filtering out the class indices [3, 4, 5], you would have to map them to [0, 1, 2].