Update app.py

app.py

@@ -1,3 +1,4 @@
+import spaces
 import gradio as gr
 import torch
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
@@ -10,6 +11,7 @@ model_name = "gpt2"
 model = GPT2LMHeadModel.from_pretrained(model_name)
 tokenizer = GPT2Tokenizer.from_pretrained(model_name)
 
+@spaces.GPU
 def get_token_probabilities(text, top_k=10):
     # Tokenize the input text
     input_ids = tokenizer.encode(text, return_tensors="pt")