Store hidden state in checkpoints

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.

Do you mean the hidden state of an RNN?
You can store tensors using torch.save.