Solved by Ravi_Raju in post #4 
              
              
                Hi, 
I actually managed to hack a workaround with the tools in torch.nn.utils.prune. The trick is save the dictionary of masks and remove the pruning modules with remove and then add them back with CustomFromMask. 
from torch.nn.utils.prune import l1_unstructured, remove, CustomFromMask,is_pruned


…