Your use case sounds like a multi-label classification, where each sample might belong to zero, one or multiple classes.
If that’s the case, you could use a final linear layer with out_features=nb_classes, use a one-hot encoded target in the same shape, and apply nn.BCEWithLogitsLoss as your loss function.