Update modeling_llava_qwen2.py

modeling_llava_qwen2.py

@@ -662,14 +662,14 @@ class LlavaMetaForCausalLM(ABC):
         return self.get_model().get_vision_tower()
 
     def encode_images(self, images):
-        image_features = self.get_model().get_vision_tower()(images)
-        image_features = self.get_model().mm_projector(image_features.to('cuda'))
+        image_features = self.get_model().get_vision_tower().cuda()(images)
+        image_features = self.get_model().mm_projector(image_features)
         return image_features
         
     def prepare_inputs_labels_for_multimodal(
             self, input_ids, position_ids, attention_mask, past_key_values, labels, images
     ):
-        vision_tower = self.get_vision_tower()
+        vision_tower = self.get_vision_tower().cuda()
         if vision_tower is None or images is None or input_ids.shape[1] == 1:
             if past_key_values is not None and vision_tower is not None and images is not None and input_ids.shape[
                 1] == 1: