Solved by ptrblck in post #5 
              
              
                Thanks for the code snippet. 
The approach using a dict might work, alternatively, you could also use: 
h = module.weight.register_hook(lambda grad, gradient_mask=gradient_mask: grad.mul_(gradient_mask))

Your current code runs into a known Python limitation of using lambda functions, since names in…