BCEwithLogitsLoss on 2d inputs

1

Hi, I have a question about the behavior of BCEWithLogitsLoss on a mini-batch of 2d inputs.

Suppose I have a Nx1xDxD output of logits and another Nx1xDxD mask, which contains 0,1,2.
Then I do loss = criterion(input, mask), does it ignore positions where value = 2?
If not, is there an argument that can ignore that value? Something similar to ignore_index
Can I do something like this that achieve the desired results?
loss = criterion(input[mask.data != 2], mask[mask.data != 2])

Thanks!