ypsoh:

Like you mentioned, because for every pruning step we’re creating a new mask and attaching it to the model, do you think this might be cause the model to grow in size incrementally and cause this issue? If this is the case, do you think it would help if I use prune.remove for every pruning step? So that the masks are basically omitted?


No, so the way it’s implemented is such that the mask keeps on getting replaced. We don’t keep on storing more and more masks with every iteration.
I’m most interested in your points 2 and 4. I don’t know exactly what torch.cuda.empty_cache() is supposed to do, but it’s not helping here, as you said. I want to understand if we can avoid setting retain_graph=True altogether and therefore potentially solve your problem (assuming that that’s what’s causing it). Are you explicitly calling .backward twice anywhere in your training loop, perhaps on two different parts of the network? If not, it might still be happening implicitly somewhere and we’d have to find out why and how. This issue is related to the mask application, I’m assuming, but I can’t yet tell in what exact way…
I’ll let you take a deeper look, but if you ever get to a super reduced version of the code that can be easily run and shows the same error, I can probably help take it apart for debugging.