If you prune with no_grad, I don’t think the mask application will be correctly applied and backpropagated through in the backward pass. This should be visible in the pruned weights getting updated, which is not the correct behavior. Worth checking.
Out of curiosity, does the issue happen for both global unstructured and random unstructured pruning?
Another dumb question: is there any recurrence in the model? Any fancy weight sharing? I’m asking because if a parameter is reused in different ways (for example when an LSTM is unrolled over the time-steps), then we might be having issues re-accessing the graph info after the first time. This would be a problem to pin down and address in pruning.