diff --git a/flow_field.py b/flow_field.py
index 7669eea..3b1f963 100644
--- a/flow_field.py
+++ b/flow_field.py
@@ -574,7 +574,7 @@ def flow_field(
 
     if pre_mask is not None:
       s = geom_utils.query_integral_image(
-          _integral_image(pre_mask), patch_size, step
+          _integral_image(pre_mask), patch_size, step  # pyrefly: ignore[bad-argument-type]
       )
       m = (s / np.prod(patch_size) >= max_masked)[out_sel]
       selection_mask[m] = False
@@ -582,7 +582,7 @@ def flow_field(
 
     if post_mask is not None:
       s = geom_utils.query_integral_image(
-          _integral_image(post_mask), post_patch_size, step
+          _integral_image(post_mask), post_patch_size, step  # pyrefly: ignore[bad-argument-type]
       )
       m = (s / np.prod(post_patch_size) >= max_masked)[out_sel]
       selection_mask[m] = False