@albanD, Hi, will be glad to hear your opinion about the correctness of this solution, because I find it the easiest to experiment with.
Think it should work ?



 soulitzer:

mod = nn.Linear(2, 2)
y_hat = mod(x)
y_hat.register_hook(lambda x: x*10)
y = mod(y_hat)
y.register_hook(lambda x: x/10)
y.sum().backward()

Does this resolve your issue?


In my case, I’ll register a hook on the am_loss which will divide by 10 and another hook on the input of the second path, which is the masked_image, which will multiply the grads back by 10, so I’ll get the light-blue path gradients divided by 10 and the blue as regular, will I ?
Another solution as we discussed is to register a hook for each parameter in the module, which is a little bit painful (and to maintain counters or as @soulitzer suggested with internal counters trick and e.t.c which is still isn’t nice). But the idea using many hook for modules which are in nn.Sequential isn’t likable, especially to to the fact that right now the backward_hook for nn.Sequential is broken in terms of it only works on the gradients of the last module in it and not on every one, so it means I need to loop over each module handly and register a hook which will divide by 10 the gradients.