Why do I converge only using a parameter wrapper in PyTorch?
I'm writing an implementation of Bayesian Dropout, and the following code worked after several hours of debugging:

class DropoutLayer(nn.Module):

    def __init__(self, n_in, n_out, dropout_rate):
        super(DropoutLayer,self).__init__()
        self.M = Parameter(torch.normal(torch.zeros(n_in,n_out),0.01))
        self.m = Parameter(torch.rand(n_out))

        # 1-p is used to be consistent with original dropout definition
        self.z = Bernoulli(torch.tensor([1-dropout_rate]))

        self.W = Parameter(torch.mm(
        torch.diagflat(self.z.sample(sample_shape=(n_in,))),
        self.M))



    def forward(self,x,act):
        activation = torch.mm(x,self.W) + self.m

        out = act(activation) 
        return out


However, I do not seem to understand, why the matrix multiplication in self.W needs a parameter wrapping. I would assume once I nominate self.M to be a parameter in autograd, I would not need to do it for any values which use it further down the line. Why do I need it here?

After adding this wrapper, the neural net converges without a problem.
Could you tell me why is it the case and if there is a simpler substitute for this solution which I don't like?

First of all, you don't need to pass self.M into Parameter, as the gradient does not go through it. You forward function only uses self.W and self.m (plus activation, although you really should pass it to the constructor, not forward...).

All self.M is some kind of random normal tensor created at the same time as module and it's simply initializing your self.W matrix to concrete values. So self.W is a separate tensor, it's not computationally dependent on self.M in any way.
