jmercat commited on
Commit
2aa2690
1 Parent(s): 5ae38c5

add spaces GPU decorator to try out zero-gpu

Browse files
Files changed (1) hide show
  1. app.py +2 -0
app.py CHANGED
@@ -1,3 +1,4 @@
 
1
  import gradio as gr
2
  from threading import Thread
3
  from open_lm.hf import *
@@ -23,6 +24,7 @@ def load_model(model_name):
23
  current_model = current_model.to(device)
24
  return f"Loaded model: {model_name}"
25
 
 
26
  def generate(
27
  prompt, model_choice, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0,
28
  ):
 
1
+ import spaces
2
  import gradio as gr
3
  from threading import Thread
4
  from open_lm.hf import *
 
24
  current_model = current_model.to(device)
25
  return f"Loaded model: {model_name}"
26
 
27
+ @spaces.GPU
28
  def generate(
29
  prompt, model_choice, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0,
30
  ):