Merge pull request #77 from WallarooLabs/fix_shape_panic

Return an Err on shape mismatch instead of panic

Author: nbigaouette

onnxruntime/src/error.rs

@@ -119,6 +119,14 @@ pub enum NonMatchingDimensionsError {
         /// Input dimensions defined in model
         model_input: Vec<Vec<Option<u32>>>,
     },
+    /// Inputs length from model does not match the expected input from inference call
+    #[error("Different input lengths: Expected Input: {model_input:?} vs Received Input: {inference_input:?}")]
+    InputsLength {
+        /// Input dimensions used by inference call
+        inference_input: Vec<Vec<usize>>,
+        /// Input dimensions defined in model
+        model_input: Vec<Vec<Option<u32>>>,
+    },
 }
 
 /// Error details when ONNX C API fail

onnxruntime/src/session.rs

@@ -520,10 +520,19 @@ impl<'a> Session<'a> {
                 "Different input lengths: {:?} vs {:?}",
                 self.inputs, input_arrays
             );
-            panic!(
-                "Different input lengths: {:?} vs {:?}",
-                self.inputs, input_arrays
-            );
+            return Err(OrtError::NonMatchingDimensions(
+                NonMatchingDimensionsError::InputsLength {
+                    inference_input: input_arrays
+                        .iter()
+                        .map(|input_array| input_array.shape().to_vec())
+                        .collect(),
+                    model_input: self
+                        .inputs
+                        .iter()
+                        .map(|input| input.dimensions.clone())
+                        .collect(),
+                },
+            ));
         }
 
         // Verify shape of each individual inputs
@@ -540,10 +549,19 @@ impl<'a> Session<'a> {
                 "Different input lengths: {:?} vs {:?}",
                 self.inputs, input_arrays
             );
-            panic!(
-                "Different input lengths: {:?} vs {:?}",
-                self.inputs, input_arrays
-            );
+            return Err(OrtError::NonMatchingDimensions(
+                NonMatchingDimensionsError::InputsLength {
+                    inference_input: input_arrays
+                        .iter()
+                        .map(|input_array| input_array.shape().to_vec())
+                        .collect(),
+                    model_input: self
+                        .inputs
+                        .iter()
+                        .map(|input| input.dimensions.clone())
+                        .collect(),
+                },
+            ));
         }
 
         Ok(())