Update app.py

app.py

@@ -94,7 +94,7 @@ def prompt_to_layout(user_prompt, intensity, fpath=None):
         
     top_p, top_k = creativity(intensity)
     model_prompt = '[User prompt] {} [Layout]'.format(user_prompt)
-    input_ids = tokenizer(model_prompt, return_tensors='pt')
+    input_ids = tokenizer(model_prompt, return_tensors='pt').to(device)
     output = finetuned.generate(**input_ids, do_sample=True, top_p=top_p, top_k=top_k, 
                                 eos_token_id=50256, max_length=400)
     output = tokenizer.batch_decode(output, skip_special_tokens=True)