Update app.py

app.py

@@ -2,11 +2,16 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
+device = torch.device('cuda')
+
 model_name = "Salesforce/codegen-350M-mono"
 
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
+model.to(device)
+tokenizer.to(device)
+
 def generate_code(prompt):
     full_prompt = f"# HTML/CSS/JS website code for:\n# {prompt}\n"
     inputs = tokenizer(full_prompt, return_tensors="pt").input_ids