A simple trick I might suggest is to reshape your inputs to (batch_size * num_sentences, max_words, embed_dim), run them through your LSTM, and then you’ll get an output of shape (batch_size * num_sentences, hidden_size) (by taking the last hidden state of the PyTorch nn.LSTM). Then you can easily reshape your matrix to be (batch_size, num_sentences, hidden_size). In this case, you have the last hidden state for each sentence, which is what you want. However, that is the case if your sentences are not connected to each other (such as paragraphs where the sentence depends on the previous one/recursion). If this is your case, I would highly recommend using a Hierarchical LSTM. Look at the PyTorch implementation here 54 for HAN implementation.