Skip to content

Commit ee7464f

Browse files
committed
[Test] correctly handle multi-channel input ranges in change_batchsize
1 parent 3c870d6 commit ee7464f

File tree

1 file changed

+5
-0
lines changed

1 file changed

+5
-0
lines changed

tests/transformation/test_change_batchsize.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,11 @@ def test_change_batchsize(test_model):
4545
batch_size = 10
4646
old_ishape = test_details["input_shape"]
4747
imin, imax = test_details["input_range"]
48+
# some models spec per-channel ranges, be conservative for those
49+
if isinstance(imin, np.ndarray):
50+
imin = imin.max()
51+
if isinstance(imax, np.ndarray):
52+
imax = imax.min()
4853
model = download_model(test_model=test_model, do_cleanup=True, return_modelwrapper=True)
4954
iname = model.graph.input[0].name
5055
oname = model.graph.output[0].name

0 commit comments

Comments
 (0)