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)