Fix external weights relative path and loop import (#619)

Signed-off-by: Kevin Chen <kevinch@nvidia.com>
onnx · Jan 11, 2021 · 21f62e9 · 21f62e9
1 parent 47e2b0f
commit 21f62e9
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/builtin_op_importers.cpp b/builtin_op_importers.cpp
@@ -1703,6 +1703,7 @@ DEFINE_BUILTIN_OP_IMPORTER(Loop)
     const ::ONNX_NAMESPACE::GraphProto& body = attrs.get<const ::ONNX_NAMESPACE::GraphProto&>("body");
 
     auto loop = ctx->network()->addLoop();
+    loop->setName(getNodeName(node).c_str());
     // Trip count and condition are optional inputs.
     nvinfer1::ITensor* tripLimit{nullptr};
     if (inputs[0])
@@ -1732,7 +1733,6 @@ DEFINE_BUILTIN_OP_IMPORTER(Loop)
         ctx->loopTensors()[body.input(i).name()] = node.input(i);
         ctx->registerTensor(TensorOrWeights{stateVars.back()->getOutput(0)}, body.input(i).name());
     }
-    ctx->registerLayer(stateVars.at(0), node.name());
 
     // Loop body
     TRT_CHECK(onnx2trt::parseGraph(ctx, body));

diff --git a/onnx2trt_utils.cpp b/onnx2trt_utils.cpp
@@ -1325,6 +1325,10 @@ bool parseExternalWeights(IImporterContext* ctx, std::string file, std::string p
     {
         path.replace(slash + 1, path.size() - (slash + 1), file);
     }
+    else
+    {
+        path = file;
+    }
     std::ifstream relPathFile(path, std::ios::binary | std::ios::ate);
     if (!relPathFile)
     {