use get_n

bioimage-io · Nov 12, 2024 · 8cd6c85 · 8cd6c85
1 parent 21c934e
commit 8cd6c85
Showing 1 changed file with 3 additions and 9 deletions.
diff --git a/bioimageio/spec/model/v0_5.py b/bioimageio/spec/model/v0_5.py
@@ -2549,17 +2549,11 @@ def get_axis_size(a: Union[InputAxis, OutputAxis]):
                         + f" '{a.id}' of tensor '{t_descr.id}'."
                     )
                 n = ns[(t_descr.id, a.id)]
-                s_n = a.size.get_size(n)
                 s_max = max_input_shape.get((t_descr.id, a.id))
                 if s_max is None:
-                    return s_n
-
-                for n_min in range(n):
-                    s = a.size.get_size(n_min)
-                    if s >= s_max:
-                        return s
-
-                return s_n  # n == 0
+                    return n
+                else:
+                    return min(n, a.size.get_n(s_max))
 
             elif isinstance(a.size, SizeReference):
                 if (t_descr.id, a.id) in ns: