Update app.py

app.py

@@ -280,8 +280,8 @@ def expand_prompt(prompt):
         input_text_2 = f"{system_prompt_rewrite} {user_prompt_rewrite_2} {prompt}"
         print("-- got prompt --")
         # Encode the input text and include the attention mask
-        encoded_inputs = txt_tokenizer(input_text, return_tensors="pt", return_attention_mask=True)
-        encoded_inputs_2 = txt_tokenizer(input_text_2, return_tensors="pt", return_attention_mask=True)
+        encoded_inputs = txt_tokenizer(input_text, return_tensors="pt", return_attention_mask=True).to("cuda:0")
+        encoded_inputs_2 = txt_tokenizer(input_text_2, return_tensors="pt", return_attention_mask=True).to("cuda:0")
         # Ensure all values are on the correct device
         input_ids = encoded_inputs["input_ids"].to("cuda:0")
         input_ids_2 = encoded_inputs_2["input_ids"].to("cuda:0")