Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import torch
+import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from transformers import StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
 from threading import Thread
@@ -20,6 +21,7 @@ class StopOnTokens(StoppingCriteria):
 
 
 # Function to generate model predictions.
+@spaces.GPU
 def predict(message, history):
     history_transformer_format = history + [[message, ""]]
     stop = StopOnTokens()