generate: missing to in DoLa body, causing exceptions in multi-gpu …

…generation (#32856)
huggingface · Aug 17, 2024 · 52cb403 · 52cb403
1 parent 6806d33
commit 52cb403
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py
@@ -2379,7 +2379,7 @@ def _dola_decoding(
             for candidate_premature_layer in candidate_premature_layers:
                 candidate_premature_logits[candidate_premature_layer] = lm_head(
                     outputs.hidden_states[candidate_premature_layer][:, -1, :]
-                )
+                ).to(final_logits.device)
 
             if synced_gpus and this_peer_finished:
                 continue  # don't waste resources running the code we don't need