My mistake, you also have to set the correct new batch size.
model.batch_size = test_batch_size
model.hidden_state = model.init_hidden()

If you don’t know why you need to do that then you know little about how an LSTM works.



 kaushalshetty:

I have a LSTM model trained for a batch size = 512. This means that 512 hidden states are initialized for each sample in the batch.


Wrong. Using a batch size of 512 simply means that you feed the model with 512 samples in parallel. It is faster to do 512 calculations in one go, than to do 512 calculations one after the other. There are hidden_dim hidden states initialised per sample.
Once you have trained the model using a batch size of 512, the resulting trained model is valid for any batch size. You just have to reinitialise the hidden state for the new batch size.
The hidden state is the model’s “short-term memory”, so each time you give your model a batch of data, then unless the new batch of data contains the continuation of the sequences in the previous batch, then you should reset the hidden state.