Update README.md

README.md

@@ -69,7 +69,7 @@ prompt = f"""<code>{code}</code>
 ```python
 from transformers import AutoModelForCausalLM, AutoTokenizer
 device = "cuda"
-model = AutoModelForCausalLM.from_pretrained("PipableAI/pip-code-to-doc-1.3b")
+model = AutoModelForCausalLM.from_pretrained("PipableAI/pip-code-to-doc-1.3b").to(device)
 tokenizer = AutoTokenizer.from_pretrained("PipableAI/pip-code-to-doc-1.3b")
 
 inputs = tokenizer(text, return_tensors="pt")