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