Update app.py

app.py

@@ -2,6 +2,7 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import numpy as np
 import gradio as gr
+import spaces
 
 tokenizer = AutoTokenizer.from_pretrained("gpt2")
 model = AutoModelForCausalLM.from_pretrained("gpt2")
@@ -382,7 +383,7 @@ def generate_beams(start_sentence, scores, length_penalty, decoded_sequences):
 
     return original_tree
 
-
+@spaces.GPU
 def get_beam_search_html(input_text, number_steps, number_beams, length_penalty):
     inputs = tokenizer([input_text], return_tensors="pt")