Update README.md

README.md

@@ -102,7 +102,9 @@ inputs = processor(
 )
 
 
-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)