myownskyW7 commited on
Commit
3b3bafa
1 Parent(s): 1543e50

Update modeling_InternLM.py

Browse files
Files changed (1) hide show
  1. modeling_InternLM.py +1 -1
modeling_InternLM.py CHANGED
@@ -1241,6 +1241,6 @@ class InternLMForCausalLM(InternLMPreTrainedModel):
1241
  reordered_past = ()
1242
  for layer_past in past_key_values:
1243
  reordered_past += (tuple(
1244
- past_state.index_select(0, beam_idx)
1245
  for past_state in layer_past), )
1246
  return reordered_past
 
1241
  reordered_past = ()
1242
  for layer_past in past_key_values:
1243
  reordered_past += (tuple(
1244
+ past_state.index_select(0, beam_idx.to(past_state.device))
1245
  for past_state in layer_past), )
1246
  return reordered_past