Ignoring a specific label during training in segmentation

Hello, I started working with deep learning with Pytorch not very long ago and I have an issue with segmentation. My labels have three values : ‘0: no data’, ‘1: no vegetation’ and ‘2: vegetation’. I don’t want my algorithm to learn and classify ‘0: no data’ labels. After digging on internet I saw that one possible solution consists in putting a weight of 0 in the loss function, but I’m thinking that it might be problematic with respect to the computation of the gradient, I don’t know if I’m right though. Is there an elegant way to ignore the ‘0: no data’ label from training apart from setting a weight to “0” in the loss function? Thank you in advance.

You could alternatively use ignore_index in nn.CrossEntropyLoss to ignore this class during the loss calculation.