trying to use zerogpu

app.py

@@ -2,6 +2,7 @@ import json
 import os
 import shutil
 import requests
+import spaces
 
 import gradio as gr
 from huggingface_hub import Repository
@@ -17,14 +18,6 @@ device = "cuda"
 tokenizer = AutoTokenizer.from_pretrained(checkpoint)
 model = AutoModelForCausalLM.from_pretrained(checkpoint, trust_remote_code=True).to(device)
 
-prompt = '<fim_prefix>def print_hello_world():\n    """<fim_suffix>\n    print("Hello world!")<fim_middle>'
-
-inputs = tokenizer.encode(prompt, return_tensors="pt").to(device)
-outputs = model.generate(inputs, max_length=100, temperature=0.2)
-#print("-"*80)
-print(tokenizer.decode(outputs[0]))
-
-
 FIM_PREFIX = "<fim_prefix>"
 FIM_MIDDLE = "<fim_middle>"
 FIM_SUFFIX = "<fim_suffix>"
@@ -46,7 +39,7 @@ theme = gr.themes.Monochrome(
 
 inputs = tokenizer.encode(prompt, return_tensors="pt").to(device)
 outputs = model.generate(inputs, max_length=100, temperature=0.2)
-
+@spaces.GPU
 def generate(
     prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0, version="StarCoder",
 ):