Fix T5 beam search using parallelize (#11717)

huggingface · May 14, 2021 · bd3b599 · bd3b599
1 parent 218d552
commit bd3b599
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/transformers/models/t5/modeling_t5.py b/src/transformers/models/t5/modeling_t5.py
@@ -1682,7 +1682,7 @@ def _reorder_cache(self, past, beam_idx):
             for layer_past_state in layer_past_states:
                 # need to set correct `past` for each of the four key / value states
                 reordered_layer_past_states = reordered_layer_past_states + (
-                    layer_past_state.index_select(0, beam_idx),
+                    layer_past_state.index_select(0, beam_idx.to(layer_past_state.device)),
                 )
 
             assert reordered_layer_past_states[0].shape == layer_past_states[0].shape