Update README.md

README.md

@@ -37,7 +37,7 @@ import json
 
 device = "cuda" # the device to load the model onto
 
-model = AutoModelForCausalLM.from_pretrained("fireworks-ai/firefunction-v1")
+model = AutoModelForCausalLM.from_pretrained("fireworks-ai/firefunction-v1", device_map="auto")
 tokenizer = AutoTokenizer.from_pretrained("fireworks-ai/firefunction-v1")
 
 function_spec = [
@@ -87,7 +87,7 @@ messages = [
     {'role': 'user', 'content': 'Hi, can you tell me the current stock price of AAPL?'}
 ]
 
-model_inputs = tokenizer.apply_chat_template(messages, return_tensors="pt")
+model_inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device)
 
 generated_ids = model.generate(model_inputs, max_new_tokens=1000, do_sample=True)
 decoded = tokenizer.batch_decode(generated_ids)