If I understand your problem. You only want to backpropagate through only those pixels that are True in the mask. In that case, you cannot use the approach, you defined in the question as that would backpropagate through all the pixels.
Use this code instead
class myModel(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(1, 10)
def forward(self, x, trgt):
x = self.linear(x)
# Mask code
mask = trgt > -1
x[mask==0] = 0
return x
model = myModel()
# Dummy inputs:- Assume mask is 10 units
x = torch.rand(1, 1)
trgt = torch.arange(-4, 6).view(1, 10).to(dtype=torch.float)
# Initial param valuies
print('Before update')
print(model.state_dict())
# Run a epoch step
# Large lr to visualize easily
optimizer = optim.SGD(model.parameters(), lr=100)
optimizer.zero_grad()
output = model(x, trgt)
# This step is copied from your question as you take diff[mask]
trgt[trgt<=-1]=0
loss = torch.mean(torch.abs(output - trgt))
loss.backward()
optimizer.step()
print('\nAfter update')
print(model.state_dict())
Before update
OrderedDict([('linear.weight', tensor([[ 0.2001],
[ 0.8265],
[ 0.4066],
[ 0.8853],
[ 0.3886],
[-0.1819],
[ 0.4687],
[-0.8151],
[ 0.6224],
[-0.1107]])), ('linear.bias', tensor([ 0.1837, 0.0316, -0.2824, 0.9075, 0.6819, 0.6282, -0.9894, -0.2376,
0.2664, -0.9535]))])
After update
OrderedDict([('linear.weight', tensor([[ 0.2001],
[ 0.8265],
[ 0.4066],
[ 0.8853],
[-7.4897],
[ 7.6964],
[ 8.3469],
[ 7.0632],
[ 8.5007],
[ 7.7676]])), ('linear.bias', tensor([ 0.1837, 0.0316, -0.2824, 0.9075, -9.3181, 10.6282, 9.0106, 9.7624,
10.2664, 9.0465]))])
As you see only those parameters gor upgraded that were part of the mask.