Hi,
I have a complex module where I want to have gradient back prop over a specific region of input. Can I do this in the forward function on the input?
def forward(self, input):
    self.loss = nn.MSELoss(input, self.target)
   
    def hook_fn(grad):
         return grad * self.mask // where mask is a binary mask defined to be a subset of input 
    input.register_hook(hook_fn)
   
    return input