Gradient flow of MSE loss after multiplication with mask

I want the loss function like ||Mask · (output − target )||2. Should the below code works fine for gradient flow? Note: Mask is bool(True or False).


import torch.nn as nn
import torch.nn.functional as F
import torch
from torchsummary import summary

class TwoLayerNet(nn.Module):
    def __init__(self):
        super(TwoLayerNet, self).__init__()
        conv_layers=[]
        self.conv=[]      

        conv_layers.append(nn.Sequential(nn.Conv2d(3, 5, kernel_size=3, stride=1,padding=1, bias=True) ) )
        conv_layers.append(nn.Sequential(nn.Conv2d(5, 5, kernel_size=3, stride=1,padding=1, bias= True),nn.LeakyReLU(0.2,True) ))
        self.conv = nn.Sequential(*conv_layers)     
                     
    def forward(self, x):
        x1=self.conv(x)
        return x1

torch.manual_seed(88)
x = torch.ones(1,3,2,2)
tar=torch.ones(1,5,2,2)*0.1

model = TwoLayerNet()
y_pred = model(x)
mask=y_pred>tar
print("mask",mask)
crit=torch.nn.MSELoss()
loss= mask*crit(y_pred,tar)

I’m facing the same error, when I add requires_grad=True it works but with no improvement at all the loss is the same for all steps no changes at all

Just check to follow optimizer.step() after backward(). It should work. My case is different. It multiply with mask to calculate loss function.

check here to get an idea about the problem I’m facing I think it’s similar to yours