code: def _update_states(self,
                         final_states: RnnStateStorage,
                         restoration_indices: torch.LongTensor) -> None:
          new_unsorted_states = [state.index_select(1, restoration_indices)
                                 for state in final_states]
          if self._states is None:
              self._states = tuple(state.data for state in new_unsorted_states)
          else:
              current_state_batch_size = self._states[0].size(1)
              new_state_batch_size = final_states[0].size(1)
              used_new_rows_mask = [(state[0, :, :].sum(-1)
                                     != 0.0).float().view(1, new_state_batch_size, 1)
                                    for state in new_unsorted_states]
              new_states = []
              if current_state_batch_size > new_state_batch_size:
                  for old_state, new_state, used_mask in zip(self._states,
                                                             new_unsorted_states,
                                                             used_new_rows_mask):
                      masked_old_state = old_state[:, :new_state_batch_size, :] * (1 - used_mask)
                      old_state[:, :new_state_batch_size, :] = new_state + masked_old_state
                      new_states.append(old_state.detach())
              else:
                  new_states = []
                  for old_state, new_state, used_mask in zip(self._states,
                                                             new_unsorted_states,
                                                             used_new_rows_mask):
                      masked_old_state = old_state * (1 - used_mask)
                      new_state += masked_old_state
                      new_states.append(new_state.detach())
              self._states = tuple(new_states)