From b7020e1cd07d998e5b086655537ac936c32d6818 Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Sun, 19 Feb 2023 17:36:54 -0500 Subject: [PATCH] Make input_shapes into a list of lists if only a list was passed --- src/deepsparse/utils/onnx.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/deepsparse/utils/onnx.py b/src/deepsparse/utils/onnx.py index d234990c32..8d6f4d98a8 100644 --- a/src/deepsparse/utils/onnx.py +++ b/src/deepsparse/utils/onnx.py @@ -236,7 +236,8 @@ def override_onnx_input_shapes( ] # Input shapes should be a list of lists, even if there is only one input - assert all(isinstance(inp, list) for inp in input_shapes) + if not all(isinstance(inp, list) for inp in input_shapes): + input_shapes = [input_shapes] # If there is a single input shape given and multiple inputs, # duplicate for all inputs to apply the same shape