Thanks for the reply.


Yes it is a relatively big model and yes I think the size of the task itself may have something to do with the error. (I’ve tried the same idea to a much lighter transformer architecture and didn’t come across any OOM errors).


As for the retain_graph=True part, I know it’s really a hack at this point but I’m applying torch.cuda.empty_cache() right after I use the retain_graph=True flag. To be honest, I’m not sure if this API is allowing the model to free the buffers is was clinging onto due to retain_graph=True.


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?


As you suggested, I’ll take a deeper look into why the trainer is complaining about calling backward twice when pruning is applied. If cuda.empty_cache is not doing what I think it’s doing then probably retain_graph=True is causing the issue.


Thank you so much. Love your work. You rock.