Oh, sorry, by bad.
LSTM returns (all_outputs, (last_hidden, last_cell) (my example),
but I believe GRU returns (all_outputs, last_hidden)
In that case using
_, last_states_packed = self.rnn...
should work