Update README.md

README.md

@@ -64,7 +64,7 @@ def _inference(tokenizer, model, generation_config, prompt, pixel_values=None):
     if pixel_values is None:
         output_tensors = model.generate(**inputs, generation_config=generation_config)
     else:
-        output_tensors = model.generate(**inputs, generation_config=generation_config, pixel_values=pixel_values.to(model.dtype))
+        output_tensors = model.generate(**inputs, generation_config=generation_config, pixel_values=pixel_values.to(model.dtype).to(model.device))
     output_str = tokenizer.decode(output_tensors[0])
     return output_str
 ```