Hello,
how would one effectively mask the parameters of a module without losing neither their link to the optimizer?
More in detail:
Approach 1:
for weights in model.parameters():
    backups.append(weights.clone().detach().data)
    mask = sample_mask(some_arguments...)  # this comes from weights of a wrapper module that needs to be trained jointly
    weights.data *= mask
loss = loss_func(model(data), target)
for backup, weights in zip(backups, model.parameters()):
    weights.data = backup

This seemed good, but the weights of model lose the grad_fn attribute that the mask accumulated for its generation process and that I need for my training!
Approach 2:
for weights in model.parameters():
    backup.append(weights.clone().detach().data)
    mask = sample_mask(some_arguments...)
    weights.data *= mask
    weights.grad_fn = mask.grad_fn  # ERROR, grad_fn is not writable!
loss = loss_func(model(data), target)

Doesn’t work as pytorch forbids assigning to grad_fn (why? It may not be common or best practice, but certainly has its use cases). If it did work, it would be my preferred approach i guess, as it makes clear what I need as well.
Approach 3:
for weights in model.parameters():
    del weights from optimizer
    mask = sample_mask(some_arguments...)
    weights = weights * mask
    optimizer.add_weights(weights) in the right place somehow
loss = loss_func(model(data), target)

Basically very hacky and unclear as to how to delete only the exact weights that I want to see deleted, and then afterwards assign them in the correct spot again…
So since approach 3 appears to be pretty unthinkable for a bigger composite model, I am a bit at a loss at how to simply mask parameters and still train the parameters and the mask generation process jointly. Tips would be appreciated