First of all, thanks for the awesome explanation of gradInput and gradOutput. I was trying to figure out what they were, but couldn’t find it anywhere in the docs.
I’ve found a workaround for this. It’s not an elegant solution, but it works.
def saliency_map_general(model, input, label):
    if (isinstance(model, torchvision.models.Inception3)):
        input = preprocess_inception(input)
    else:
        input = preprocess(input)
    
    h = [0]*len(list(model.modules()))

    def hookfunc(module, gradInput, gradOutput):
        print('hook callback is running')
        // do something here
    for j, i in enumerate(list(model.modules())):
        h[j] = i.register_backward_hook(hookfunc)
    output = model.forward(input)
    model.zero_grad()
                                        
    output[0][label].backward()
                                        
    for i in range(len(list(model.modules()))):
        h[i].remove()
    grads = input.grad.data.clamp(min=0)
    grads.squeeze_()
    grads.transpose_(0,1)
    grads.transpose_(1,2)
    grads = np.amax(grads.cpu().numpy(), axis=2)
    
    return grads

Can someone guide me as to how I should clamp the gradients to 0? (where I wrote “// do something here” in the code).
gradInput and gradOutput are both tuples, so I’m not able to do any operations on them.
I cannot convert gradInput to type “list” because it takes way too much time to do that (I grew impatient and restarted the jupyter notebook after a few minutes).