Browse Source

TRT 4 compatibility

John Welsh 6 năm trước cách đây
mục cha
commit
de66de0be2
1 tập tin đã thay đổi với 3 bổ sung1 xóa
  1. 3 1
      src/uff_to_plan.cpp

+ 3 - 1
src/uff_to_plan.cpp

@@ -68,7 +68,7 @@ int main(int argc, char *argv[])
   IBuilder *builder = createInferBuilder(gLogger);
   INetworkDefinition *network = builder->createNetwork();
   IUffParser *parser = createUffParser();
-  parser->registerInput(inputName.c_str(), DimsCHW(3, inputHeight, inputWidth));
+  parser->registerInput(inputName.c_str(), DimsCHW(3, inputHeight, inputWidth), UffInputOrder::kNHWC);
   parser->registerOutput(outputName.c_str());
   if (!parser->parse(uffFilename.c_str(), *network, dataType))
   {
@@ -81,7 +81,9 @@ int main(int argc, char *argv[])
 
   /* build engine */
   if (dataType == DataType::kHALF)
+  {
     builder->setHalf2Mode(true);
+  }
 
   builder->setMaxBatchSize(maxBatchSize);
   builder->setMaxWorkspaceSize(maxWorkspaceSize);