Solved by justusschock in post #2 
              
              
                You could simply write an initialization function, which could look like this: 
def reset_parameters(self):
    for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.B…