I think you are right, the code I have here prunes the weight gradients after the backpropagation. Contributions from pruned weights are still included, which is incorrect. Well, it is a hacky attempt I had - to do it better we need to create masks for each layer using register_register_hook I think.
Any thought?