Fix T5 beam search using parallelize (#11717)

bd3b599c · Oyvind Tafjord · GitHub · 218d552f · bd3b599c
Unverified Commit bd3b599c authored May 14, 2021 by Oyvind Tafjord Committed by GitHub May 14, 2021
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

src/transformers/models/t5/modeling_t5.py src/transformers/models/t5/modeling_t5.py +1 -1

No files found.
--- a/src/transformers/models/t5/modeling_t5.py
+++ b/src/transformers/models/t5/modeling_t5.py
@@ -1682,7 +1682,7 @@ class T5ForConditionalGeneration(T5PreTrainedModel):
            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