So imagine at some point in the training I want to multiply my weights by a mask before they’re used in the forward method (I want to freeze the weights and learn the mask instead).
I can of course inherit all the layers I’m interested in and add this mask and override the forward method but I was thinking if I could dynamically remove the weight parameter in the modules and monkey-patch them with something like below:
class CustomWeight(nn.Module):
    def __init__(self, weight):
        self.weight_mask = nn.Parameter(torch.ones_like(weight))
        
        self.weight = weight.clone()
        self.weight.requires_grad = False

    def forward(self, x):
        return self.weight * self.weight_mask

Does this make sense?