Update app.py

app.py

@@ -1,11 +1,15 @@
 import gradio as gr
 from transformers import pipeline
+import torch
 
-# Load the model with optimization settings
+# Check if a GPU is available
+device = 0 if torch.cuda.is_available() else -1
+
+# Load the text-generation pipeline with the appropriate device
 model = pipeline(
     "text-generation",
     model="rish13/polymers",
-    device=0  # Use device=0 for GPU (if available), -1 for CPU
+    device=device  # Automatically use GPU if available, otherwise CPU
 )
 
 def generate_response(prompt):