gpu util

app.py

@@ -1,17 +1,15 @@
 import gradio as gr
 from transformers import pipeline
-import torch
-
-# Check if a GPU is available
-device = 0 if torch.cuda.is_available() else -1
-print("Using GPU" if device == 0 else "Using CPU")
-
-# Load the model on the GPU if available
-model = pipeline("text-generation", model="gpt2", device=device)
+from spaces import GPU  # Import the GPU decorator for ZeroGPU
 
+# Decorate the function to indicate it needs GPU resources
+@GPU
 def generate_text(prompt):
+    # Load the model within the function so that it only runs on GPU when the function is called
+    model = pipeline("text-generation", model="gpt2", device=0)
     return model(prompt, max_length=50)[0]["generated_text"]
 
+# Create the Gradio interface
 interface = gr.Interface(
     fn=generate_text,
     inputs=gr.Textbox(label="Enter your prompt here"),