Update app.py

app.py

@@ -1,5 +1,11 @@
 import gradio as gr
 
+torch.set_default_device("cuda")
+
+model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", torch_dtype="auto", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)
+
+
 def generate_primes(text_input):
     inputs = tokenizer(text_input, return_tensors="pt", return_attention_mask=False)
     outputs = model.generate(**inputs, max_length=200)