adapt test to multi tensor model

FynnBe · FynnBe · commit f155c9dcb08d · 2024-11-15T15:57:40.000+01:00
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(

Original file line number	Diff line number	Diff line change
`@@ -355,7 +355,7 @@ def split_multiple_shapes_into_blocks(`
`355`	`355`	`halo=halo.get(t, {}),`
`356`	`356`	`stride=strides.get(t),`
`357`	`357`	`)`
`358`		`- assert n_blocks[t] > 0`
	`358`	`+ assert n_blocks[t] > 0, n_blocks`
`359`	`359`
`360`	`360`	`assert len(blocks) > 0, blocks`
`361`	`361`	`assert len(n_blocks) > 0, n_blocks`