RuntimeError: running_mean should contain 1 elements not 256

I’m trying to train an simple encoder decoder network, while training it is working proper but while loading back the model and running inference, I’m getting following error.

 RuntimeError: running_mean should contain 1 elements not 256

Encoder Network


class LSTMEncoder(nn.Module):
    """_summary_

    Args:
        nn (_type_): _description_
    """

    def __init__(self, input_size, hidden_size=265):
        super(LSTMEncoder, self).__init__()

        self.embed = nn.Embedding(input_size, hidden_size,)

        self.rnn_unit = nn.LSTM(hidden_size, hidden_size,
                                dropout=0.2, num_layers=2, batch_first=True)
        self.rnn_unit2 = nn.LSTM(hidden_size, hidden_size,
                                dropout=0.2, num_layers=2, batch_first=True)
        
        
        self.batch_norm = nn.BatchNorm1d(hidden_size)
        self.dropout = nn.Dropout(0.2)

    def forward(self, x):
        """_summary_

        Args:
            x (_type_): _description_

        Returns:
            _type_: _description_
        """
        x = self.embed(x)
        x = relu(x)
        out, (h_n, c_n) = self.rnn_unit(x)
        out = self.batch_norm(out) --> At this point error ioccurring
        out, (h_n, c_n) = self.rnn_unit2(out)
        
        return out, h_n, c_n