Skip to content

Commit e9a0748

Browse files
committed
Fix external weights relative path and loop import (#619)
Signed-off-by: Kevin Chen <kevinch@nvidia.com>
1 parent b297dda commit e9a0748

File tree

2 files changed

+5
-1
lines changed

2 files changed

+5
-1
lines changed

builtin_op_importers.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1703,6 +1703,7 @@ DEFINE_BUILTIN_OP_IMPORTER(Loop)
17031703
const ::ONNX_NAMESPACE::GraphProto& body = attrs.get<const ::ONNX_NAMESPACE::GraphProto&>("body");
17041704

17051705
auto loop = ctx->network()->addLoop();
1706+
loop->setName(getNodeName(node).c_str());
17061707
// Trip count and condition are optional inputs.
17071708
nvinfer1::ITensor* tripLimit{nullptr};
17081709
if (inputs[0])
@@ -1732,7 +1733,6 @@ DEFINE_BUILTIN_OP_IMPORTER(Loop)
17321733
ctx->loopTensors()[body.input(i).name()] = node.input(i);
17331734
ctx->registerTensor(TensorOrWeights{stateVars.back()->getOutput(0)}, body.input(i).name());
17341735
}
1735-
ctx->registerLayer(stateVars.at(0), node.name());
17361736

17371737
// Loop body
17381738
TRT_CHECK(onnx2trt::parseGraph(ctx, body));

onnx2trt_utils.cpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1325,6 +1325,10 @@ bool parseExternalWeights(IImporterContext* ctx, std::string file, std::string p
13251325
{
13261326
path.replace(slash + 1, path.size() - (slash + 1), file);
13271327
}
1328+
else
1329+
{
1330+
path = file;
1331+
}
13281332
std::ifstream relPathFile(path, std::ios::binary | std::ios::ate);
13291333
if (!relPathFile)
13301334
{

0 commit comments

Comments
 (0)