Given a base model, I want to implement another model which recomputes the weights every time (W = W0 * func(W1) ) and uses W as the module weights for computations, but trains only the parameters W0 and W1 (instead of training W).
I have been trying the following, but this doesn’t seem to work:
class AdvModel(nn.Module):
    def __init__(self, model):
        super(AdvModel, self).__init__()
        self.model = model
        for param in select_params(model): # iterates on selected parameters of model
            # (W) :this need not be updated
            param.weight.requires_grad = False 
            # (W0) this will be updated by autograd
            param.register_parameter(name='W0', param=deepcopy(param.weight) ) 
            # (W1) Two cases:
            #  this is updated by autograd; or 
            #  this is updated manually after each iteration
            param.register_parameter(name='W1', 
                param=torch.nn.Parameter(torch.rand(param.weight.shape), 
                requires_grad=False))  
            
    def forward(self, x):
        for param in select_params(self.model):
            adv_W1 = func_1(param.W1) 
            # where func_1 maybe a non-differentiable/differentiable function
            param.weight = torch.nn.Parameter(param.W0 * adv_W1, 
                                             requires_grad=False)
        return self.model(x)