use GPU

app.py

@@ -16,10 +16,10 @@ example = [
     [EXAMPLE_3, 11, 0.2, 42],
     ]
 
-# change model to the finetuned one
+device="cuda:0"
 tokenizer = AutoTokenizer.from_pretrained("loubnabnl/santacoder-code-to-text")
-model = AutoModelForCausalLM.from_pretrained("loubnabnl/santacoder-code-to-text", trust_remote_code=True)
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+model = AutoModelForCausalLM.from_pretrained("loubnabnl/santacoder-code-to-text", trust_remote_code=True).to(device)
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)
 
 def make_doctring(gen_prompt):
     return gen_prompt + f"\n\n\"\"\"\nExplanation:"