Update app.py

app.py

@@ -22,12 +22,17 @@ def main():
             description = generate_description(title)
         st.success(description)
 
+# Check if GPU is available
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+else:
+    device = torch.device("cpu")
 
 # Define the function that generates the description
 def generate_description(title):
     # Preprocess the input
     input_text = f"{title}"
-    input_ids = tokenizer.encode(input_text, return_tensors='pt')
+    input_ids = tokenizer.encode(input_text, return_tensors='pt').to(device)
 
     # Generate the output using the model
     output = model.generate(