Update modeling_qwen2_vl.py

huggingface · Sep 4, 2024 · bb4c2b6 · bb4c2b6
1 parent 2fbec02
commit bb4c2b6
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py b/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py
@@ -1579,16 +1579,12 @@ def forward(
                 pixel_values = pixel_values.type(self.visual.get_dtype())
                 image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw).to(inputs_embeds.device)
                 image_mask = input_ids == self.config.image_token_id
-                if self.training:
-                    inputs_embeds = inputs_embeds.clone()
-                inputs_embeds[image_mask] = image_embeds
+                inputs_embeds.masked_scatter(image_mask.unsqueeze(-1), image_embeds)
             if pixel_values_videos is not None:
                 pixel_values_videos = pixel_values_videos.type(self.visual.get_dtype())
                 video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw).to(inputs_embeds.device)
                 video_mask = input_ids == self.config.video_token_id
-                if self.training:
-                    inputs_embeds = inputs_embeds.clone()
-                inputs_embeds[video_mask] = video_embeds
+                inputs_embeds.masked_scatter(video_mask.unsqueeze(-1), video_embeds)
             if attention_mask is not None:
                 attention_mask = attention_mask.to(inputs_embeds.device)