Update app.py

app.py

@@ -73,7 +73,7 @@ def encode_prompt(prompt_batch, text_encoder, tokenizer, proportion_empty_prompt
             max_length=256,
             truncation=True,
             return_tensors="pt",
-        )
+        ).cuda()
 
         print(f"Text Encoder Device: {text_encoder.device}")
         text_input_ids = text_inputs.input_ids.cuda()