I’m doing the following – please note that I support both GRU and LSTM in the model so I can decide on setup time:
def forward(self, inputs):
    batch_size = inputs.shape[0]
    # Push through embedding layer
    X = self.embedding(inputs)
    # Push through RNN layer (the ouput is irrelevant)
    _, self.hidden = self.rnn(X, self.hidden)

    # Get the hidden state of the last layer of the RNN
    if self.params.rnn_type == RnnType.RNN_TYPE__GRU:
        hidden = self.hidden
    elif self.params.rnn_type == RnnType.RNN_TYPE__LSTM:
        hidden = self.hidden[0]
    # Flatten hidden state with respect to batch size
    hidden = hidden.transpose(1,0).contiguous().view(batch_size, -1)
    ...

The important part is the transpose(1,0) to get the batch size to the  front. Everything else is handled by the view() command. I only need the transpose since I initialize the RNN with batch_first=True.
Note that the input shape of the directly following linear layer needs to be (rnn_hidden_dim * num_directions * num_layers, output_size).