For a Linear module with parameters W and b, grad_i will be a tuple that includes gradient over input, W and b. So should we instead use the following function?

module.register_backward_hook(lambda module, grad_i, grad_o: (grad_i[0] * -1, grad_i[1], grad_i[2]))