I am trying to use the nn.BatchNorm1d in my linear layer. The problem is that the code for neural network creation is not as simple as the examples given for batchnorm1d implementaion and is part of bigger reinforcement learning program (taken from datahubbs) and the code is made in such a way to be flexible w.r.t number of neurons and number of layers.
class actorCriticNet(nn.Module):
    def __init__(self, learning_rate, n_hidden_layers, n_hidden_nodes,
                  bias=False, device='cpu'):
        super(actorCriticNet, self).__init__()
        
        self.device = device
        self.n_inputs = 23
        self.n_outputs = 11
        self.n_hidden_nodes = n_hidden_nodes
        self.n_hidden_layers = n_hidden_layers
        self.learning_rate = learning_rate
        self.bias = bias
        self.action_space = np.arange(self.n_outputs)
        
        # Generate network according to hidden layer and node settings
        self.layers = OrderedDict()
        self.n_layers = 2 * self.n_hidden_layers
        for i in range(self.n_layers + 1):
            # Define single linear layer
            if self.n_hidden_layers == 0:
                self.layers[str(i)] = nn.Linear(
                    self.n_inputs,
                    self.n_outputs,
                    bias=self.bias)
            # Define input layer for multi-layer network
            elif i % 2 == 0 and i == 0 and self.n_hidden_layers != 0:
                self.layers[str(i)] = nn.Linear( 
                    self.n_inputs, 
                    self.n_hidden_nodes,
                    bias=self.bias)
            # Define intermediate hidden layers
            elif i % 2 == 0 and i != 0:
                self.layers[str(i)] = nn.Linear(
                    self.n_hidden_nodes,
                    self.n_hidden_nodes,
                    bias=self.bias)
            else:
                self.layers[str(i)] = nn.ReLU(nn.BatchNorm1d(self.n_hidden_nodes))
                
        self.body = nn.Sequential(self.layers)
            
        # Define policy head
        self.policy = nn.Sequential(
            nn.Linear(self.n_hidden_nodes,
                      self.n_hidden_nodes,
                      bias=self.bias),
            nn.ReLU(),
            nn.Linear(self.n_hidden_nodes,
                      self.n_outputs,
                      bias=self.bias))
        # Define value head
        self.value = nn.Sequential(
            nn.Linear(self.n_hidden_nodes,
                      self.n_hidden_nodes,
                      bias=self.bias),
            nn.ReLU(),
            nn.Linear(self.n_hidden_nodes,
                      1, 
                      bias=self.bias))

        
        self.optimizer = torch.optim.Adam(self.parameters(),
                                          lr=self.learning_rate)
    
      
    def predict(self, state):
        body_output = self.get_body_output(state)
        probs = F.softmax(self.policy(body_output), dim=-1)
        return probs, self.value(body_output)

    def get_body_output(self, state):
        state_t = torch.FloatTensor(state).to(device=self.device)
        return self.body(state_t)
    
    def get_action(self, state):
        probs = self.predict(state)[0][0][0].detach().numpy()
        action = np.random.choice(self.action_space, p=probs)
        return action
    
    def get_log_probs(self, state):
        body_output = self.get_body_output(state)
        logprobs = F.log_softmax(self.policy(body_output), dim=-1)
        return logprobs

The line of interest being:
self.layers[str(i)] = nn.ReLU(nn.BatchNorm1d(self.n_hidden_nodes))
Is my code for the implementation of batchnorm1d correct? or should I modify the code?
I am newbie in python/pytorch so please point out if there are any obvious optimisation mistakes I have made.