|
@@ -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);
|