Update app.py

app.py

@@ -138,7 +138,7 @@ def prompt_with_mutation(user_prompt, top_p, top_k, mut_rate, fpath=None):
     rooms[-1] = rooms[-1].split(':')[0] + ':'
     rooms = ', '.join(rooms)# + ', '
     new_prompt = '[User prompt] {} [Layout] {}'.format(user_prompt, rooms)
-    input_ids = tokenizer(new_prompt, return_tensors='pt')
+    input_ids = tokenizer(new_prompt, return_tensors='pt').to(device)
     mut_output = finetuned.generate(**input_ids, do_sample=True, top_p=top_p, top_k=top_k, eos_token_id=50256, max_length=400)
     mut_output = tokenizer.batch_decode(mut_output, skip_special_tokens=True)
     mut_im = mut_txt2layout(mut_output)