Update README.md

README.md

@@ -168,7 +168,8 @@ import requests
 from transformers import AutoModelForCausalLM, AutoProcessor
 
 # Load the model and processor
-model = AutoModelForCausalLM.from_pretrained("microsoft/Magma-8B", trust_remote_code=True)
+dtype = torch.bfloat16
+model = AutoModelForCausalLM.from_pretrained("microsoft/Magma-8B", trust_remote_code=True, torch_dtype=dtype)
 processor = AutoProcessor.from_pretrained("microsoft/Magma-8B", trust_remote_code=True)
 model.to("cuda")
 
@@ -185,7 +186,7 @@ prompt = processor.tokenizer.apply_chat_template(convs, tokenize=False, add_gene
 inputs = processor(images=[image], texts=prompt, return_tensors="pt")
 inputs['pixel_values'] = inputs['pixel_values'].unsqueeze(0)
 inputs['image_sizes'] = inputs['image_sizes'].unsqueeze(0)
-inputs = inputs.to("cuda")
+inputs = inputs.to("cuda").to(dtype)
 
 generation_args = { 
     "max_new_tokens": 128,