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


def add_masks(model,masks):
    mask_pruner = CustomFromMask(None)
    for module_name, module in model.named_modules():
        key = f"{module_name}.weight_mask"
        if key in masks:
            if isinstance(module, torch.nn.Conv2d):
                _mask = masks[key]
                mask_pruner.apply(module, 'weight', _mask)
            if isinstance(module, torch.nn.Linear):
                _mask = masks[key]
                mask_pruner.apply(module, 'weight', _mask)

def merge_masks(model):
    for n, m in model.named_modules():
        if is_pruned(m)==False:
            continue
        if isinstance(m, torch.nn.Conv2d):
            remove(m, name='weight')
        if isinstance(m, torch.nn.Linear):
            remove(m, name='weight')

With this fix, I can reinitialize my network with no issues. Hopefully this can help someone if they were dealing with my issue.