diff --git a/processor/flow.py b/processor/flow.py
index b9dd342..dfcaeec 100644
--- a/processor/flow.py
+++ b/processor/flow.py
@@ -210,6 +210,6 @@ def _estimate_flow(z_prev, z_curr):
 
   def expected_output_box(
       self, box: bounding_box.BoundingBoxBase) -> bounding_box.BoundingBoxBase:
-    scale_factor = 1 / self.pixelsize(np.array([1, 1, 1]))
+    scale_factor = 1 / self.pixelsize(np.repeat(1, len(box.size)))
     cropped_box = self.crop_box(box)
     return cropped_box.scale(list(scale_factor))