fix device map (#132)

* fix device map

Author: hjh0119

internvl_chat/internvl/model/internlm2/modeling_internlm2.py

@@ -1082,13 +1082,16 @@ def forward(
             output = (logits,) + outputs[1:]
             return (loss,) + output if loss is not None else output
 
-        return CausalLMOutputWithPast(
+        device = input_ids.device if input_ids is not None else inputs_embeds.device
+        output = CausalLMOutputWithPast(
             loss=loss,
             logits=logits,
             past_key_values=outputs.past_key_values,
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
         )
+        output['logits'] = output['logits'].to(device)
+        return output
 
     def prepare_inputs_for_generation(
         self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs

internvl_chat/internvl/model/internvl_chat/modeling_internvl_chat.py

@@ -226,7 +226,7 @@ def extract_feature(self, pixel_values):
         vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
         vit_embeds = self.pixel_shuffle(vit_embeds, scale_factor=self.downsample_ratio)
         vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1])
-        vit_embeds = self.mlp1(vit_embeds)
+        vit_embeds = self.mlp1(vit_embeds).to(pixel_values.device)
         return vit_embeds
 
     def chat(self, tokenizer, pixel_values, question, generation_config, history=None, return_history=False,