Hi, I don’t know if this is the best solution, but maybe you can pass a mask as an input to the forward function:
def forward(self, x, mask):
    # do something
    return mask * outputs

And than use the same mask on the targets:
masked_outputs = model(inputs, mask)
masked_targets = mask * targets
loss = criterion(masked_outputs, masked_targets)

I haven’t tried this is PyTorch yet, please tell me if it works for you!