I want to change the gradients during a backward pass for each Conv2d modules, so I’m trying to figure out how to change the input_gradiens using the backward hook, but cannot figure out what to return from the hook function in order to change the input_gradients.

def backward_hook(module, input_grads, output_grads):
        if isinstance(module, nn.Conv2d):
            input_grads = ..... # changing the input_grads here
            return input_grads, output_grads

    for module in model.modules():
        module.register_backward_hook(backward_hook)

And I’m getting this error on calling the .backward()
TypeError: expected Variable, but hook returned 'tuple'

I’ve tried all permutations and combinations 🥲 such as:
return input_grads[0], output_grads[0]

return input_grads, output_grads[0]

return input_grads,

return input_grads[0]

…
etc but still no luck.
What should I return in the case of Conv2d? Also, I could not find any documentation related to this, so very difficult to get some results.
I’ve also noticed that what we need to return in the backward hook function also depends on the type of the module, for instance the code below runs just fine with the nn.ReLU:

def backward_hook(module, input_grads, output_grads):
        if isinstance(module, nn.ReLU):
            input_grads = ..... # changing the input_grads here
            return input_grads, # so with relu it seems we don't need to return output_grads, 
            # but this does not work with Conv2d

    for module in model.modules():
        module.register_backward_hook(backward_hook)