Update README.md

README.md

@@ -57,7 +57,7 @@ model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 def inference(prompt: str) -> str:
-  model_inputs = tokenizer([prompt], return_tensors="pt").to(device)
+  model_inputs = tokenizer([prompt], return_tensors="pt").to('cuda')
   generated_ids = model.generate(model_inputs.input_ids, max_new_tokens=512)
   generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)]
   response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]