Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 from transformers import pipeline
 
 # Load the TinyLlama model for text generation
-pipe = pipeline("text-generation", model="TinyLlama/TinyLlama_v1.1")
+pipe = pipeline("text-generation", model="TinyLlama/TinyLlama_v1.1").to("cuda")
 
 # Define the prediction function
 def generate_text(prompt, max_length=128, temperature=1.0, top_p=0.95):