From 1fae113c5c9469de201fdd8c34efbdf9f76c59d3 Mon Sep 17 00:00:00 2001
From: Matthew Deville <matthew@stockly.ai>
Date: Thu, 22 Jan 2026 01:28:12 +0100
Subject: [PATCH] update

---
 src/main.rs | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/main.rs b/src/main.rs
index ea8aed4..5fc4158 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -63,12 +63,14 @@ fn main() -> Result<(), Box<dyn Error>> {
     // Create ONNX Runtime session
     let mut session = create_session(&model_path)?;
 
+    let input_name = session.inputs()[0].name().to_string();
+
     // Preprocess image
     let input_tensor = preprocess_image(&img, model_info.input_size.0, model_info.input_size.1)?;
 
     // Run inference
     println!("Running inference...");
-    let outputs = session.run(ort::inputs!["input" => Tensor::from_array(input_tensor)?])?;
+    let outputs = session.run(ort::inputs![input_name => Tensor::from_array(input_tensor)?])?;
 
     // Extract mask output
     let mask_output = &outputs[0];