Update app.py

app.py

@@ -1,10 +1,12 @@
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
 import gradio as gr
+import torch
 
 trained_tokenizer = GPT2Tokenizer.from_pretrained("Kumarkishalaya/GPT-2-next-word-prediction")
 trained_model = GPT2LMHeadModel.from_pretrained("Kumarkishalaya/GPT-2-next-word-prediction")
 untrained_model = GPT2Tokenizer.from_pretrained("gpt2")
 untrained_tokenizer = ("gpt2")
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def generate(commentary_text):
     input_ids = trained_tokenizer(commentary_text, return_tensors="pt")
@@ -13,7 +15,7 @@ def generate(commentary_text):
     return tokenizer_finetuned.decode(output[0])
 
 # Create Gradio interface
-iface = gr.Interface(fn=generate_text, 
+iface = gr.Interface(fn=generate, 
                      inputs="text", 
                      outputs="text",
                      title="GPT-2 Text Generation",