Fix zeroGPU.

app.py

@@ -4,6 +4,7 @@ import requests
 import pytz
 import yaml
 from tools.final_answer import FinalAnswerTool
+import spaces
 
 from Gradio_UI import GradioUI
 
@@ -45,7 +46,10 @@ temperature=0.5,
 model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
 custom_role_conversions=None,
 )
-
+class ZeroGPUCodeAgent(CodeAgent):
+    @spaces.GPU
+    def run(self, *args, **kwargs):
+        return super().run(*args, **kwargs)
 
 # Import tool from Hub
 image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
@@ -53,7 +57,7 @@ image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_co
 with open("prompts.yaml", 'r') as stream:
     prompt_templates = yaml.safe_load(stream)
     
-agent = CodeAgent(
+agent = ZeroGPUCodeAgent(
     model=model,
     tools=[final_answer], ## add your tools here (don't remove final answer)
     max_steps=6,