How to fix those weights for all the layers
torch.sqrt(2/n_in)
this is the more general code if i am not wrong
for m in self.modules():
    if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
        import scipy.stats as stats
        stddev = m.stddev if hasattr(m, 'stddev') else 0.1
        X = stats.truncnorm(-2, 2, scale=stddev)
        values = torch.Tensor(X.rvs(m.weight.data.numel()))
        values = values.view(m.weight.data.size())
        m.weight.data.copy_(values)
    elif isinstance(m, nn.BatchNorm2d):
        m.weight.data.fill_(1)
        m.bias.data.zero_()