Add noise to layer output

Hi, I am trying to add noise to layers’ output. I take the code from stylegan2-pytorch.

I thought the self.weight is a learnable parameter thus it can update during training since it is nn.Parameter(), but when I track the value of grad and weight of this layer the weight remains 0 and grad is None from the beginning. Can any one give me suggestions about this? thanks.

class NoiseInjection(nn.Module):
    def __init__(self):
        super().__init__()

        self.weight = nn.Parameter(torch.zeros(1), requires_grad=True)

    def forward(self, x, noise=None):
        if noise is None:
            batch, _, height, width = x.shape
            noise = x.new_empty(batch, 1, height, width).normal_()
        return x + self.weight * noise

I cannot reproduce the issue using your code snippet and get a valid gradient:

model = NoiseInjection()
x = torch.randn(16, 3, 24, 24)
out = model(x)
out.mean().backward()
print(model.weight.grad)
> tensor([0.0115])

Thanks, I double check my scripts and notice I didn’t pass the value correctly.