Update app.py

app.py

@@ -3,6 +3,8 @@ import streamlit as st
 
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+import spaces
+
 checkpoint = "."
   
 tokenizer = AutoTokenizer.from_pretrained(checkpoint)
@@ -14,6 +16,7 @@ def load_model(model_name):
 
 model = load_model(checkpoint)
 
+@spaces.GPU
 def infer(input_ids, bad_words_ids, max_tokens, temperature, top_k, top_p):
 
     output_sequences = model.generate(