Update README.md

README.md

@@ -101,7 +101,9 @@ inputs = processor(
     return_tensors="pt",
 )
 
-inputs = {k: v.to("cuda") for k, v in inputs.items() if isinstance(v, torch.Tensor)}
+inputs = {k: v.cuda() if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
+if "pixel_values" in inputs:
+    inputs["pixel_values"] = inputs["pixel_values"].to(torch.bfloat16)
 
 output_ids = model.generate(**inputs, max_new_tokens=128)
 response = processor.decode(output_ids[0], skip_special_tokens=True)