Sorry, maybe I confused you.
I need something link this:
model = nn.Conv2d(3, 1, 3, 1, 1)
mask = torch.sigmoid(model(x)) # but I need mask to be binary instead of values between 0 and 1
masked_img = mask * x # point-wise
loss = cal_loss(masked_img, x)
loss.mean().backward()