fix fuyu device_map compatibility (#29880)

fix foward

fix fuyu device_map compatibility (#29880)
fix foward
31c575bc · Marc Sun · GitHub · 4d8427f7 · 31c575bc
Unverified Commit 31c575bc authored Mar 27, 2024 by Marc Sun Committed by GitHub Mar 27, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 1 deletion

src/transformers/models/fuyu/modeling_fuyu.py src/transformers/models/fuyu/modeling_fuyu.py +3 -1

No files found.
--- a/src/transformers/models/fuyu/modeling_fuyu.py
+++ b/src/transformers/models/fuyu/modeling_fuyu.py
@@ -290,7 +290,9 @@ class FuyuForCausalLM(FuyuPreTrainedModel):
            inputs_embeds = self.language_model.get_input_embeddings()(input_ids)
            if image_patches is not None and past_key_values is None:
                patch_embeddings = [
-                    self.vision_embed_tokens(patch.to(self.vision_embed_tokens.weight.dtype)).squeeze(0)
+                    self.vision_embed_tokens(patch.to(self.vision_embed_tokens.weight.dtype))
+                    .squeeze(0)
+                    .to(inputs_embeds.device)
                    for patch in image_patches
                ]
                inputs_embeds = self.gather_continuous_embeddings(