Update app.py

app.py

@@ -10,7 +10,7 @@ a 1.1B parameter model for code generation in Python, Java & JavaScript. The mod
 with the <span style='color: #ff75b3;'>&lt;FILL-HERE&gt;</span> token.</span>"""
 
 token = os.environ["HUB_TOKEN"]
-device="cuda"
+device="cuda:0"
 
 
 FIM_PREFIX = "<fim-prefix>"
@@ -30,7 +30,7 @@ tokenizer_fim.add_special_tokens({
 
 tokenizer = AutoTokenizer.from_pretrained("bigcode/christmas-models", use_auth_token=token)
 model = AutoModelForCausalLM.from_pretrained("bigcode/christmas-models", trust_remote_code=True, use_auth_token=token).to(device)
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)
 
 def post_processing(prompt, completion):
     completion = "<span style='color: #ff75b3;'>" + completion + "</span>"