From 31e34e95eab5e3165d07847b66b41374f5fe9ebd Mon Sep 17 00:00:00 2001 From: SnowGushiGit Date: Tue, 4 Dec 2018 15:26:48 +0900 Subject: [PATCH] Update registerInput `registerInput` is changed https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/c_api/classnvuffparser_1_1_i_uff_parser.html#aa88fc7901e6805be9e093dd0b98310e3 --- src/tensorNet.cpp | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/tensorNet.cpp b/src/tensorNet.cpp index 57f2b5c..d5029de 100644 --- a/src/tensorNet.cpp +++ b/src/tensorNet.cpp @@ -56,12 +56,12 @@ ICudaEngine* createTrtFromUFF(char* modelpath) { auto parser = createUffParser(); - parser->registerInput("enc_text", DimsCHW(1, VOC_LEN, 1)); - parser->registerInput("dec_text", DimsCHW(1, VOC_LEN, 1)); - parser->registerInput("h0_in", DimsCHW(1, DIM, 1)); - parser->registerInput("c0_in", DimsCHW(1, DIM, 1)); - parser->registerInput("h1_in", DimsCHW(1, DIM, 1)); - parser->registerInput("c1_in", DimsCHW(1, DIM, 1)); + parser->registerInput("enc_text", DimsCHW(1, VOC_LEN, 1), UffInputOrder::kNCHW); + parser->registerInput("dec_text", DimsCHW(1, VOC_LEN, 1), UffInputOrder::kNCHW); + parser->registerInput("h0_in", DimsCHW(1, DIM, 1), UffInputOrder::kNCHW); + parser->registerInput("c0_in", DimsCHW(1, DIM, 1), UffInputOrder::kNCHW); + parser->registerInput("h1_in", DimsCHW(1, DIM, 1), UffInputOrder::kNCHW); + parser->registerInput("c1_in", DimsCHW(1, DIM, 1), UffInputOrder::kNCHW); parser->registerOutput("h0_out"); parser->registerOutput("c0_out");