Solved by ludwigwinkler in post #2 
              
              
                Okay, I was able to solve the problem by realizing that the register_hook command for non-leaf nodes has to be placed precisely between forward and backward pass. 
In comparison the register_forward_hook and register_backward_hook can be initialized before the forward pass. 
While it makes a lot of …