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];