Hi, I am trying to implement the lottery ticket hypothesis using torch.nn.utils.prune but I am having trouble trying to reset the model parameters. From what I understand, the pruning module adds a new parameter, weight_orig to each module by overriding weight. When I try to reinitialize the model parameters with model.apply(reinit_weights), the parameters corresponding to weight_orig in each module is not changing. Below is some code that demonstrates. Does someone have an explanation or potential workaround?
Thanks so much in advance.
# this allows the model to train with the new optimizer
# but doesn't reset parameters that have the mask added so fails the test_model_change
def reinit_weights(m):
    if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
        m.reset_parameters()

def test_model_change(prev_iter_dict, model):
    for name, param in model.named_parameters():
        prev_param = prev_iter_dict[name]
        assert not torch.allclose(prev_param,param), 'model not updating'

model = LeNet5()
optimizer = optim.Adam(model.parameters(), lr=lr)
model.train().to(device)

pruner = L1Unstructured(0.0)
# adds masks of all ones to each of the layers
for n, m in model.named_modules():
    if isinstance(m, torch.nn.Conv2d):
        pruner.apply(m, name='weight', amount=0.0)
    if isinstance(m, torch.nn.Linear):
        pruner.apply(m, name='weight', amount=0.0)

train(model,1,train_loader,device,optimizer) # train for one epoch to test
full_cap = copy.deepcopy(model.state_dict())
model.apply(reinit_weights) # this reinitalizes the weight parameter but doesn't change the new_weight parameter
test_model_change(full_cap, model) # passes for conv0.bias but fails for conv0.weight