I want to use LayerNorm with LSTM, but I’m not sure what is the best way to use them together.
My code is as follows:
rnn = nn.LSTMCell(in_channels, hidden_dim)

hidden, cell = rnn(x, (hidden, cell))

So, if I want to add LayerNorm to this model, I will do it like this?
rnn = nn.LSTMCell(in_channels, hidden_dim)
norm = nn.LayerNorm(hidden_dim)

hidden, cell = rnn(x, (hidden, cell))
hidden = norm(hidden)
cell = norm(cell)

Is it the right way to use it? I’m confused. Sorry for my noob question.