Update app.py

app.py

@@ -6,6 +6,8 @@ import gradio as gr
 #Setting device to cuda
 torch.set_default_device("cuda")
 
+device = 'cuda' if torch.cuda.is_availble() else 'cpu'
+
 
 model = "deepapaikar/katzbot-phi2"
 
@@ -23,7 +25,7 @@ def predict_answer(question, token=25):
     messages = [{"role": "user", "content": f"{question}"}]
     
     
-    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True).to(device)
     
     
     outputs = pipeline(prompt, max_new_tokens=token, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)