Make input_shapes into a list of lists if only a list was passed (#914)

neuralmagic · Feb 21, 2023 · e40b40e · e40b40e
1 parent 2bb1bdd
commit e40b40e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/deepsparse/utils/onnx.py b/src/deepsparse/utils/onnx.py
@@ -268,7 +268,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