Update app.py

app.py

@@ -3,6 +3,7 @@ import json
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
+import spaces  # Import the spaces module for ZeroGPU compatibility
 
 # Ensure the environment has access to a CUDA-capable GPU
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -55,6 +56,7 @@ def generate_synthetic_math_problems(num_problems):
     
     return problems
 
+@spaces.GPU(duration=60)
 def solve_problem(problem):
     print(f"Solving problem: {problem}")
     with torch.no_grad():
@@ -76,6 +78,7 @@ def solve_problem(problem):
     print(f"Problem: {problem}, Answer: {answer}")
     return answer
 
+@spaces.GPU(duration=120)
 def generate_and_solve_problems(num_problems):
     problems = generate_synthetic_math_problems(num_problems)
     solved_problems = []
@@ -103,4 +106,4 @@ iface = gr.Interface(
     description="Generate and solve synthetic math problems using a HuggingFace model."
 )
 
-iface.launch()
+iface.launch()