Update README.md

README.md

@@ -59,7 +59,7 @@ import torch
 
 # Load the model and processor
 model_id = "mistral-community/pixtral-12b"
-model = LlavaForConditionalGeneration.from_pretrained(model_id, torch_dtype=torch.float16, device_map="mps")
+model = LlavaForConditionalGeneration.from_pretrained(model_id, torch_dtype=torch.float16, device_map="cuda")
 processor = AutoProcessor.from_pretrained(model_id)
 
 # Load the LoRA configuration
@@ -76,7 +76,7 @@ image = Image.open(image_path)
 PROMPT = "<s>[INST]Describe the chess position in the image, piece by piece.[IMG][/INST]"
 
 # Pass single image instead of list of URLs
-inputs = processor(text=PROMPT, images=image, return_tensors="pt").to("mps")
+inputs = processor(text=PROMPT, images=image, return_tensors="pt").to("cuda")
 generate_ids = lora_model.generate(**inputs, max_new_tokens=650)
 output = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
 print(output)