Context:
Progressive Growing Of GANs For Improved Quality, Stability, And Variation
(https:// arxiv .org /pdf/1710.10196.pdf) # I'm a new user, hence can't have more than 2 links
Specifically,

Equalized_LR.PNG753×290 113 KB


I’ve found two (vaguely) similar questions here:


Changing weight after forward and before backward 16


Adjusting parameters in the forward pass 35.


Currently, this is my implementation (which I suspect to be wrong):
class WeightScaledModule(nn.Module):
    def __init__(self, module):
        super().__init__()
        self.module = module
        self.weight = self.module.weight
    def forward(self, x):
        fan = torch.nn.init._calculate_correct_fan(self.module.weight, 'fan_in')
        self.module.weight = nn.Parameter(self.module.weight * (2/fan)**.5)
        return self.module(x)

Any help on how to “dynamically” adjust parameters would be appreciated!