fix: tensorrt output layer lookup now throws when layer does not exist

jolibrain · Nov 9, 2020 · ba7c839 · ba7c839
1 parent 9edf1d4
commit ba7c839
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/backends/tensorrt/tensorrtlib.cc b/src/backends/tensorrt/tensorrtlib.cc
@@ -298,7 +298,10 @@ namespace dd
       throw MLLibInternalException("Error while parsing caffe model "
                                    "for conversion to TensorRT");
 
-    network->markOutput(*blobNameToTensor->find(out_blob.c_str()));
+    auto bloboutput = blobNameToTensor->find(out_blob.c_str());
+    if (!bloboutput)
+      throw MLLibBadParamException("Cannot find output layer " + out_blob);
+    network->markOutput(*bloboutput);
 
     if (out_blob == "detection_out")
       network->markOutput(*blobNameToTensor->find("keep_count"));