Using register_forward_hook for recurrent networks


Hi, I have a question regarding using register_forward_hook for recurrent networks. Suppose I have a module that contains a LSTMCell as a submodule which runs for N time steps in one forward() call. I want to examine the hidden states of this LSTMCell at every time step. But, if I do register_forward_hook for that LSTMCell, I think I will only get the hidden states (h_t, c_t) at the very last time step. Am I right? If so, what is a good way to store the hidden states of the LSTMCell at every time step of a single forward pass? Should I just explicitly define variables to store them in the module definition?