adapt test to multi tensor model

bioimage-io · Nov 15, 2024 · f155c9d · f155c9d
1 parent 37acc79
commit f155c9d
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/bioimageio/core/block_meta.py b/bioimageio/core/block_meta.py
@@ -355,7 +355,7 @@ def split_multiple_shapes_into_blocks(
             halo=halo.get(t, {}),
             stride=strides.get(t),
         )
-        assert n_blocks[t] > 0
+        assert n_blocks[t] > 0, n_blocks
 
     assert len(blocks) > 0, blocks
     assert len(n_blocks) > 0, n_blocks

diff --git a/tests/test_prediction.py b/tests/test_prediction.py
@@ -91,16 +91,17 @@ def test_predict_with_blocking(with_procs: bool, prep: Prep):
 
 
 def test_predict_with_fixed_blocking(prep: Prep):
-    block_along = list(prep.input_sample.members)[0]
+    block_along = list(prep.input_sample.members)
     input_block_shape: Mapping[MemberId, Mapping[AxisId, int]] = {
-        block_along: {
+        ba: {
             "x": min(  # pyright: ignore[reportAssignmentType]
-                32, prep.input_sample.members[block_along].tagged_shape[AxisId("x")]
+                128, prep.input_sample.members[ba].tagged_shape[AxisId("x")]
             ),
             AxisId("y"): min(
-                32, prep.input_sample.members[block_along].tagged_shape[AxisId("y")]
+                128, prep.input_sample.members[ba].tagged_shape[AxisId("y")]
             ),
         }
+        for ba in block_along
     }
     try:
         out = predict(