update for batch infer

modeling_InternLM.py

@@ -1,15 +1,14 @@
 import math
-from typing import List, Union
-from typing import Optional, Tuple
+from typing import List, Optional, Tuple, Union
 
 import torch
 import torch.utils.checkpoint
-import torch.utils.checkpoint
 from einops import rearrange
 from torch import nn
 from torch.nn import CrossEntropyLoss
 from transformers.activations import ACT2FN
-from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_outputs import (BaseModelOutputWithPast,
+                                           CausalLMOutputWithPast)
 from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import logging
 
@@ -863,6 +862,6 @@ class InternLMForCausalLM(InternLMPreTrainedModel):
         reordered_past = ()
         for layer_past in past_key_values:
             reordered_past += (tuple(
-                past_state.index_select(0, beam_idx)
+                past_state.index_select(0, beam_idx.to(past_state.device))
                 for past_state in layer_past), )
         return reordered_past