If you’ve passed all parameters to the optimizer after initializing the model via:
optimizer = optim.SGD(model.parameters(), lr=1e-3)

the self.mask parameter won’t be included and thus not updated, since it wasn’t initialized yet.
Could you pass the desired input shape to the initialization of the model and create self.mask in its __init__ method?
Based on your code it seems that self.mask is initialized to a static shape in the first forward pass.