Avoid assertion failure to check output is_tensor

onnx · Nov 24, 2020 · 8135b65 · 8135b65
1 parent 7eb49ff
commit 8135b65
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ModelImporter.cpp b/ModelImporter.cpp
@@ -530,8 +530,8 @@ Status ModelImporter::importModel(
     for (::ONNX_NAMESPACE::ValueInfoProto const& output : graph.output())
     {
         ASSERT(_importer_ctx.tensors().count(output.name()), ErrorCode::kINVALID_GRAPH);
-        ASSERT(_importer_ctx.tensors().at(output.name()).is_tensor(), ErrorCode::kUNSUPPORTED_GRAPH);
-        nvinfer1::ITensor* output_tensor_ptr = &_importer_ctx.tensors().at(output.name()).tensor();
+        nvinfer1::ITensor* output_tensor_ptr
+            = &convertToTensor(_importer_ctx.tensors().at(output.name()), &_importer_ctx);
         LOG_VERBOSE("Marking " << output_tensor_ptr->getName() << " as output: " << output.name());
         output_tensor_ptr->setName(output.name().c_str());
         if (output_tensor_ptr->isNetworkInput())