You can use below
def normal_init(m, mean, std):
    if isinstance(m, (nn.Linear, nn.Conv2d)):
        m.weight.data.normal_(mean, std)
        if m.bias.data is not None:
            m.bias.data.zero_()
    elif isinstance(m, (nn.BatchNorm2d, nn.BatchNorm1d)):
        m.weight.data.fill_(1)
        if m.bias.data is not None:
            m.bias.data.zero_()

in the network class
 self.weight_init()
 def weight_init(self):
        for block in self._modules:
            try:
                for m in self._modules[block]:
                    normal_init(m,mean,std)
            except:
                normal_init(block)