How can I store the hidden state of the final layers in checkpoints? I want to use the saved hidden state as the initial hidden state in inference mode.
I have tried using Embedding or Parameters， but it seems not working well.