out-of-place scatter

src/transformers/models/blip_2/modeling_blip_2.py

@@ -1774,7 +1774,8 @@ def forward(
         # otherwise we expand manually by concating
         if hasattr(self.config, "image_token_index"):
             special_image_mask = (input_ids == self.config.image_token_index).unsqueeze(-1).expand_as(inputs_embeds)
-            inputs_embeds[special_image_mask] = language_model_inputs.flatten()
+            language_model_inputs = language_model_inputs.to(inputs_embeds.device, inputs_embeds.dtype)
+            inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, language_model_inputs)
         else:
             logger.warning_once(
                 "Expanding inputs for image tokens in BLIP-2 should be done in processing. "

src/transformers/models/blip_2/processing_blip_2.py

@@ -56,8 +56,6 @@ class Blip2Processor(ProcessorMixin):
     image_processor_class = "BlipImageProcessor"
     tokenizer_class = "AutoTokenizer"
 
-
-    # Copied from transformers.models.blip.processing_blip.BlipProcessor.__init__
     def __init__(self, image_processor, tokenizer, num_query_tokens=None, **kwargs):
         tokenizer.return_token_type_ids = False
         self.current_processor = image_processor

src/transformers/models/llava/modeling_llava.py

@@ -501,7 +501,8 @@ def forward(
                 special_image_mask = (
                     (input_ids == self.config.image_token_index).unsqueeze(-1).expand_as(inputs_embeds)
                 )
-                inputs_embeds[special_image_mask] = image_features.flatten()
+                image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
+                inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
 
         outputs = self.language_model(
             attention_mask=attention_mask,

src/transformers/models/llava/processing_llava.py

@@ -45,10 +45,10 @@ class LlavaProcessor(ProcessorMixin):
         vision_feature_select_strategy (`str`, *optional*):
             The feature selection strategy used to select the vision feature from the vision backbone.
             Shoudl be same as in model's config
-        image_token (`str`, *optional*, defaults to `"<image>"`):
-            Special token used to denote image location.
         chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages
             in a chat into a tokenizable string.
+        image_token (`str`, *optional*, defaults to `"<image>"`):
+            Special token used to denote image location.
     """
 
     attributes = ["image_processor", "tokenizer"]

src/transformers/models/llava_next/modeling_llava_next.py

@@ -867,7 +867,8 @@ def forward(
                 special_image_mask = (
                     (input_ids == self.config.image_token_index).unsqueeze(-1).expand_as(inputs_embeds)
                 )
-                inputs_embeds[special_image_mask] = image_features.flatten()
+                image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
+                inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
 
         outputs = self.language_model(
             attention_mask=attention_mask,

src/transformers/models/llava_next/processing_llava_next.py

@@ -46,10 +46,10 @@ class LlavaNextProcessor(ProcessorMixin):
         vision_feature_select_strategy (`str`, *optional*):
             The feature selection strategy used to select the vision feature from the vision backbone.
             Shoudl be same as in model's config
-        image_token (`str`, *optional*, defaults to `"<image>"`):
-            Special token used to denote image location.
         chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages
             in a chat into a tokenizable string.
+        image_token (`str`, *optional*, defaults to `"<image>"`):
+            Special token used to denote image location.
     """
 
     attributes = ["image_processor", "tokenizer"]

src/transformers/models/paligemma/modeling_paligemma.py

@@ -414,7 +414,8 @@ def forward(
                     f"Got {image_tokens_in_text} image tokens in the text but {image_features.shape[0] * image_features.shape[1]} "
                     "tokens from image embeddings."
                 )
-            inputs_embeds[special_image_mask] = image_features.flatten()
+            image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
+            inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
 
         causal_mask = self._update_causal_mask(
             attention_mask, token_type_ids, inputs_embeds, cache_position, is_training

src/transformers/models/video_llava/modeling_video_llava.py

@@ -614,13 +614,15 @@ def forward(
                     special_image_mask = (
                         (input_ids == self.config.image_token_index).unsqueeze(-1).expand_as(inputs_embeds)
                     )
-                    inputs_embeds[special_image_mask] = image_features.flatten()
+                    image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
+                    inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
 
                 if video_outputs is not None:
                     special_image_mask = (
                         (input_ids == self.config.video_token_index).unsqueeze(-1).expand_as(inputs_embeds)
                     )
-                    inputs_embeds[special_image_mask] = video_features.flatten()
+                    video_features = video_features.to(inputs_embeds.device, inputs_embeds.dtype)
+                    inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, video_features)
 
         outputs = self.language_model(
             attention_mask=attention_mask,

src/transformers/models/vipllava/modeling_vipllava.py

@@ -498,7 +498,8 @@ def forward(
                 special_image_mask = (
                     (input_ids == self.config.image_token_index).unsqueeze(-1).expand_as(inputs_embeds)
                 )
-                inputs_embeds[special_image_mask] = image_features.flatten()
+                image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
+                inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
 
         outputs = self.language_model(
             attention_mask=attention_mask,