diff --git a/tmva/sofie/test/input_models/Linear_16.onnx b/tmva/sofie/test/input_models/Linear_16.onnx
index 8e65190bba494..22f9cb95bfb70 100644
Binary files a/tmva/sofie/test/input_models/Linear_16.onnx and b/tmva/sofie/test/input_models/Linear_16.onnx differ
diff --git a/tutorials/machine_learning/TMVA_SOFIE_ONNX.C b/tutorials/machine_learning/TMVA_SOFIE_ONNX.C
index 8c192789e1210..34def293e8834 100644
--- a/tutorials/machine_learning/TMVA_SOFIE_ONNX.C
+++ b/tutorials/machine_learning/TMVA_SOFIE_ONNX.C
@@ -7,44 +7,41 @@
 /// \macro_code
 /// \macro_output
 /// \author Sanjiban Sengupta
-
 using namespace TMVA::Experimental;
-
 void TMVA_SOFIE_ONNX(std::string inputFile = ""){
    if (inputFile.empty() )
       inputFile = std::string(gROOT->GetTutorialsDir()) + "/machine_learning/Linear_16.onnx";
-
     //Creating parser object to parse ONNX files
     SOFIE::RModelParser_ONNX parser;
     SOFIE::RModel model = parser.Parse(inputFile, true);
-
     //Generating inference code
     model.Generate();
-    // write the code in a file (by default Linear_16.hxx and Linear_16.dat
+    // write the code in a file (by default Linear_16.hxx and Linear_16.dat)
     model.OutputGenerated();
-
     //Printing required input tensors
     model.PrintRequiredInputTensors();
-
     //Printing initialized tensors (weights)
     std::cout<<"\n\n";
     model.PrintInitializedTensors();
-
     //Printing intermediate tensors
     std::cout<<"\n\n";
     model.PrintIntermediateTensors();
-
     //Checking if tensor already exist in model
-    std::cout<<"\n\nTensor \"16weight\" already exist: "<<std::boolalpha<<model.CheckIfTensorAlreadyExist("16weight")<<"\n\n";
-    std::vector<size_t> tensorShape = model.GetTensorShape("16weight");
-    std::cout<<"Shape of tensor \"16weight\": ";
-    for(auto& it:tensorShape){
-        std::cout<<it<<",";
+    const std::string tensorName = "0weight";
+    bool tensorExists = model.CheckIfTensorAlreadyExist(tensorName);
+    std::cout<<"\n\nTensor \""<<tensorName<<"\" already exist: "<<std::boolalpha<<tensorExists<<"\n\n";
+    if (tensorExists) {
+        std::vector<size_t> tensorShape = model.GetTensorShape(tensorName);
+        std::cout<<"Shape of tensor \""<<tensorName<<"\": ";
+        for(auto& it:tensorShape){
+            std::cout<<it<<",";
+        }
+        std::cout<<"\n\nData type of tensor \""<<tensorName<<"\": ";
+        SOFIE::ETensorType tensorType = model.GetTensorType(tensorName);
+        std::cout<<SOFIE::ConvertTypeToString(tensorType);
+    } else {
+        std::cout<<"Tensor \""<<tensorName<<"\" not found in model. Skipping shape/type queries.\n";
     }
-    std::cout<<"\n\nData type of tensor \"16weight\": ";
-    SOFIE::ETensorType tensorType = model.GetTensorType("16weight");
-    std::cout<<SOFIE::ConvertTypeToString(tensorType);
-
     //Printing generated inference code
     std::cout<<"\n\n";
     model.PrintGenerated();