Fixed beam search error when using multiple GPUs (#6)

- Fixed beam search error when using multiple GPUs (14a911cdcbe50c9fbaae49af6f26b3d981b30cde)


Co-authored-by: Hiroki Yamaguchi <[email protected]>

modeling_plamo.py

@@ -701,5 +701,5 @@ class PlamoForCausalLM(PlamoPreTrainedModel):
     def _reorder_cache(past_key_values: List[torch.FloatTensor], beam_idx: int) -> Tuple[Any, ...]:
         reordered_past: Tuple[Any, ...] = ()
         for layer_past in past_key_values:
-            reordered_past += (tuple(past_state.index_select(0, beam_idx) for past_state in layer_past),)
+            reordered_past += (tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),)
         return reordered_past