Update app.py

app.py

@@ -15,6 +15,10 @@ from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM
 finetuned = AutoModelForCausalLM.from_pretrained('model')
 tokenizer = AutoTokenizer.from_pretrained('gpt2')
 
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+print(device)
+finetuned = finetuned.to(device)
+
 # Utility functions
 
 housegan_labels = {"living_room": 1, "kitchen": 2, "bedroom": 3, "bathroom": 4, "missing": 5, "closet": 6, 
@@ -66,7 +70,7 @@ def prompt_to_layout(user_prompt, top_p, top_k, fpath=None):
     
     model_prompt = '[User prompt] {} [Layout]'.format(user_prompt)
     #print(model_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)
     #print(output)