Solved by albanD in post #2 
              
              
                Hi, 
IIRC there is no need for backward hook as the pruning is done in a differentiable manner during the forward and so the autograd will give the right gradients.