diff --git a/README.md b/README.md
index 10a61461..0506b8cb 100644
--- a/README.md
+++ b/README.md
@@ -124,6 +124,7 @@ You can also add custom tools to the agent:
 
 ```python
 import vision_agent as va
+import numpy as np
 
 @va.tools.register_tool(imports=["import numpy as np"])
 def custom_tool(image_path: str) -> str:
@@ -140,7 +141,6 @@ def custom_tool(image_path: str) -> str:
     >>> custom_tool("image.jpg")
     """
 
-    import numpy as np
     return np.zeros((10, 10))
 ```
 
diff --git a/examples/custom_tools/template_match.py b/examples/custom_tools/template_match.py
index bb0ca416..65da51e8 100644
--- a/examples/custom_tools/template_match.py
+++ b/examples/custom_tools/template_match.py
@@ -54,6 +54,11 @@ def template_matching_with_rotation(
     for angle in range(0, max_rotation, step):
         # Rotate the template
         rotated_template = rotate_image(template_gray, angle)
+        if (
+            rotated_template.shape[0] > main_image_gray.shape[0]
+            or rotated_template.shape[1] > main_image_gray.shape[1]
+        ):
+            continue
 
         # Perform template matching
         result = cv2.matchTemplate(
@@ -69,17 +74,18 @@ def template_matching_with_rotation(
             )
             scores.append(result[y, x])
 
-    indices = (
-        nms(
-            torch.tensor(boxes).float(),
-            torch.tensor(scores).float(),
-            0.2,
+    if len(boxes) > 0:
+        indices = (
+            nms(
+                torch.tensor(boxes).float(),
+                torch.tensor(scores).float(),
+                0.2,
+            )
+            .numpy()
+            .tolist()
         )
-        .numpy()
-        .tolist()
-    )
-    boxes = [boxes[i] for i in indices]
-    scores = [scores[i] for i in indices]
+        boxes = [boxes[i] for i in indices]
+        scores = [scores[i] for i in indices]
 
     if visualize:
         # Draw a rectangle around the best match