Update README.md

README.md

@@ -21,12 +21,12 @@ import torch
 from transformers import pipeline, AutoTokenizer
 
 question = "Write a Python function that takes a list of numbers and returns the list sorted in ascending order without using the built-in `sorted()` function."
-generator = pipeline("text-generation", model="mariasandu/python-coding-assistant-v2", device="cpu")
+generator = pipeline("text-generation", model="mariasandu/python-coding-assistant-v2", device="cuda")
 
 tokenizer = AutoTokenizer.from_pretrained("mariasandu/python-coding-assistant-v2")
 tokenizer.chat_template = "{% for message in messages %}{% if message['role'] == 'user' %}[INST] {{ message['content'] }} [/INST]{% else %}{{ message['content'] }}{% endif %}{% endfor %}"
-output = generator(question)
-print(output["generated_text"])
+output = generator(question)[0]
+print(output['generated_text'])
 ```
 
 ## Training procedure