Update app.py

app.py

@@ -3,6 +3,7 @@ import os
 import shutil
 import requests
 import spaces
+import torch
 
 import gradio as gr
 from huggingface_hub import Repository
@@ -11,6 +12,8 @@ from peft import PeftModel
 
 from share_btn import community_icon_html, loading_icon_html, share_js, share_btn_css
 
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 CHECKPOINT_URL = "Salesforce/codegen-350M-mono"
 
@@ -143,8 +146,9 @@ theme = gr.themes.Monochrome(
     ],
 )
 
-
+@spaces.GPU(enable_queue=True)
 def stream(model, code, generate_kwargs):
+    model.to(device)
     input_ids = tokenizer(code, return_tensors="pt").to("cuda")
     generated_ids = model.generate(**input_ids, **generate_kwargs)
     return tokenizer.decode(generated_ids[0][input_ids["input_ids"].shape[1]:], skip_special_tokens=True).strip()