docstr: This function is used to re-order the :obj:`past_key_values` cache if
  :meth:`~transformers.PreTrainedModel.beam_search` or :meth:`~transformers.PreTrainedModel.beam_sample` is
  called. This is required to match :obj:`past_key_values` with the correct beam_idx at every generation step.
docstr_tok: ["This","function","is","used","to","re-order","the",":","obj",":","`","past_key_values","`","cache","if",":","meth",":","`","~transformers.PreTrainedModel.beam_search","`","or",":","meth",":","`","~transformers.PreTrainedModel.beam_sample","`","is","called",".","This","is","required","to","match",":","obj",":","`","past_key_values","`","with","the","correct","beam_idx","at","every","generation","step","."]
code: def _reorder_cache(
      past: Tuple[Tuple["torch.Tensor"]], beam_idx: "torch.Tensor"
  ) -> Tuple[Tuple["torch.Tensor"]]:
      """
      This function is used to re-order the :obj:`past_key_values` cache if
      :meth:`~transformers.PreTrainedModel.beam_search` or :meth:`~transformers.PreTrainedModel.beam_sample` is
      called. This is required to match :obj:`past_key_values` with the correct beam_idx at every generation step.
      """
      return tuple(
          tuple(
              past_state.index_select(0, beam_idx.to(past_state.device))
              for past_state in layer_past
          )
          for layer_past in past
      )
  
imports: from dataclasses import dataclass
  from typing import Optional, Tuple
  from .configuration_gpt2 import GPT2Config
  from fastNLP.transformers.torch.activations import ACT2FN
  from fastNLP.transformers.torch.file_utils import ModelOutput, add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, replace_return_docstrings
  from fastNLP.transformers.torch.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions, CausalLMOutputWithCrossAttentions, SequenceClassifierOutputWithPast, TokenClassifierOutput
  from fastNLP.transformers.torch.modeling_utils import Conv1D, PreTrainedModel, SequenceSummary, find_pruneable_heads_and_indices, prune_conv1d_layer
  from fastNLP.transformers.torch.utils.model_parallel_utils import assert_device_map, get_device_map
  from fastNLP.envs.imports import _NEED_IMPORT_TORCH
  from fastNLP.core.log import logger
label_desc: Use image node_id 8 for calling the GPT2DoubleHeadsModel obj's underlying member method code with example usage: obj._reorder_cache(past, beam_idx) and returns: tuple