On another note, I believe pruning the gradients using your code will result in the wrong update. That’s because gradients depend on each other. So if you zero some gradients for some parameters in some layer, then the gradients in the earlier layers should change (and might not necessarily be zero) due to the chain rule. Your code does not change the dependent gradients. So the code is probably slow and will result in the wrong update.