Update app.py

app.py

@@ -12,7 +12,10 @@ prompt = st.text_area("Enter a prompt:", "this news is real pyresearch given rig
 
 # Load model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", torch_dtype="auto", device=device, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", torch_dtype="auto", trust_remote_code=True)
+
+# Move the model to the desired device
+model.to(device)
 
 # Generate text on button click
 if st.button("Generate"):