Hook function with more parameters

You don’t need the extra signatures. You can use up-values to store it for you. Something like

parameter1 = ...
def hook(module, grad_input, grad_output):
    return grad_input, grad_output, parameter1

You can even make it better, by slightly modifying the code in Why cant I see .grad of an intermediate variable?