Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import torch
 import pandas as pd
 import plotly.graph_objects as go
+import spaces
 from plotly.subplots import make_subplots
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import time
@@ -73,7 +74,8 @@ def truncate_context(input_ids, max_length=2048):
     if len(input_ids[0]) > max_length:
         return input_ids[:, -max_length:]
     return input_ids
-
+    
+@spaces.GPU
 def generate_response(message, chat_history, epsilon):
     global data, stop_generation
     data = pd.DataFrame(columns=["Time taken (in ms)", "Early exit depth", "Token"])