diff --git a/vision_agent/agent/vision_agent.py b/vision_agent/agent/vision_agent.py
index ef0f305c..7d937e3b 100644
--- a/vision_agent/agent/vision_agent.py
+++ b/vision_agent/agent/vision_agent.py
@@ -368,7 +368,8 @@ def _handle_viz_tools(
         image_to_data[image]["scores"].extend(call_result.get("scores", []))
         image_to_data[image]["masks"].extend(call_result.get("masks", []))
         # only single heatmap is returned
-        image_to_data[image]["heat_map"].append(call_result.get("heat_map", []))
+        if "heat_map" in call_result:
+            image_to_data[image]["heat_map"].append(call_result["heat_map"])
         if "mask_shape" in call_result:
             image_to_data[image]["mask_shape"] = call_result["mask_shape"]