I made a module that uses the following MLP module:
class MLP(nn.Module):
    def __init__(self, size_layers, activation):
        super(MLP, self).__init__()
        self.layers=[]
        self.layersnorm = []
        self.activation=activation
        for i in range(len(size_layers)-1):
            self.layers.append(nn.Linear(size_layers[i], size_layers[i + 1]))
            self.add_module('layers_' + str(i),self.layers[-1])

            self.layersnorm.append(nn.BatchNorm1d(size_layers[i + 1]))
            self.add_module('BatchNorm1d_' + str(i), self.layersnorm[-1])

    def forward(self, x):
        for i in range(len(self.layers)-1):
            if self.activation=='relu':
                x = F.relu(self.layersnorm[i](self.layers[i](x)))
            elif self.activation=='lrelu':
                x = F.leaky_relu(self.layersnorm[i](self.layers[i](x)))
            elif self.activation=='tanh':
                x = F.tanh(self.layersnorm[i](self.layers[i](x)))
        x = self.layersnorm[-1](self.layers[-1](x))
        return x
    def l1reg(self):
        w=0.
        for i in range(len(self.layers)):
            w = w + torch.sum((self.layers[i].weight).abs())
        return w

everything works fine without batch normalization.
With batch normalization the training seem to work, but the evaluation (using model.eval()) produces NaN.
is there something I’m doing wrong with batch normalization?
thanks!