code: def _get_initial_states(self,
                              batch_size: int,
                              num_valid: int,
                              sorting_indices: torch.LongTensor) -> Optional[RnnState]:
          if self._states is None:
              return None
          if batch_size > self._states[0].size(1):
              num_states_to_concat = batch_size - self._states[0].size(1)
              resized_states = []
              for state in self._states:
                  zeros = state.new_zeros(state.size(0),
                                          num_states_to_concat,
                                          state.size(2))
                  resized_states.append(torch.cat([state, zeros], 1))
              self._states = tuple(resized_states)
              correctly_shaped_states = self._states
          elif batch_size < self._states[0].size(1):
              correctly_shaped_states = tuple(state[:, :batch_size, :] for state in self._states)
          else:
              correctly_shaped_states = self._states
          if len(self._states) == 1:
              correctly_shaped_state = correctly_shaped_states[0]
              sorted_state = correctly_shaped_state.index_select(1, sorting_indices)
              return sorted_state[:, :num_valid, :].contiguous()
          else:
              sorted_states = [state.index_select(1, sorting_indices)
                               for state in correctly_shaped_states]
              return tuple(state[:, :num_valid, :].contiguous() for state in sorted_states)