Update README.md

README.md

@@ -32,7 +32,7 @@ tokenizer.pad_token = tokenizer.eos_token
 tokenizer.chat_template = AutoTokenizer.from_pretrained(CHAT_TEMPLATE_ID).chat_template
 
 model = AutoModelForCausalLM.from_pretrained(
-        model_name, device=device, trust_remote_code=True, dtype=torch.float16)
+        model_name, device_map=device, trust_remote_code=True)
 
 messages = []
 prompt = "Tell me 5 sites to visit in Spain"