diff --git a/core/conversion/conversion.cpp b/core/conversion/conversion.cpp index bebcf794c8..27287fe568 100644 --- a/core/conversion/conversion.cpp +++ b/core/conversion/conversion.cpp @@ -1,6 +1,6 @@ -#include - #include "core/conversion/conversion.h" +#include +#include #include "core/conversion/conversionctx/ConversionCtx.h" #include "core/conversion/converters/converters.h" #include "core/conversion/evaluators/evaluators.h" @@ -234,10 +234,28 @@ void MarkOutputs(ConversionCtx* ctx, at::ArrayRef outp } } } else { - std::string name = std::string("output_") + std::to_string(ctx->num_outputs); + bool setOutput = false; + auto num_inputs = ctx->net->getNbInputs(); auto out_tensor = it->second; - out_tensor->setName(name.c_str()); - ctx->net->markOutput(*out_tensor); + std::string name = std::string("output_") + std::to_string(ctx->num_outputs); + + // Check if the output tensor is one of the inputs to the network. If so, apply an identity layer to it. + for (int64_t i = 0; i < num_inputs; i++) { + if (out_tensor == ctx->net->getInput(i)) { + LOG_DEBUG( + "One of the inputs named " + << ctx->net->getInput(i)->getName() + << " to the network is marked as an output tensor. Applying an identity layer and marking this tensor as output"); + auto id_out_tensor = converters::applyIdentityOp(ctx, out_tensor, name); + ctx->net->markOutput(*id_out_tensor); + setOutput = true; + } + } + + if (!setOutput) { + out_tensor->setName(name.c_str()); + ctx->net->markOutput(*out_tensor); + } LOG_INFO( ctx->logger, "Marking Output " << out->debugName() << " named " << name << " in engine (ctx.MarkOutput)"); ctx->num_outputs += 1; diff --git a/core/conversion/converters/converter_util.cpp b/core/conversion/converters/converter_util.cpp index 146ce1b80d..9312706b47 100644 --- a/core/conversion/converters/converter_util.cpp +++ b/core/conversion/converters/converter_util.cpp @@ -121,6 +121,13 @@ nvinfer1::ILayer* add_elementwise( return ele; } +nvinfer1::ITensor* applyIdentityOp(ConversionCtx* ctx, nvinfer1::ITensor* tensor, const std::string& tensor_name) { + auto id_layer = ctx->net->addIdentity(*tensor); + auto id_out_tensor = id_layer->getOutput(0); + id_out_tensor->setName(tensor_name.c_str()); + return id_out_tensor; +} + nvinfer1::ITensor* castITensor(ConversionCtx* ctx, nvinfer1::ITensor* tensor, nvinfer1::DataType dtype) { if (tensor->getType() != dtype) { std::ostringstream tensor_id; diff --git a/core/conversion/converters/converter_util.h b/core/conversion/converters/converter_util.h index 03808ed8ce..3f134666d6 100644 --- a/core/conversion/converters/converter_util.h +++ b/core/conversion/converters/converter_util.h @@ -41,6 +41,9 @@ nvinfer1::ILayer* add_elementwise( nvinfer1::ITensor* other, const std::string& name); +// Apply an identity operation on a tensor. Used in the case where an input is an output to a network. +nvinfer1::ITensor* applyIdentityOp(ConversionCtx* ctx, nvinfer1::ITensor* tensor, const std::string& name); + // If an ITensor is of a type not dtype, add an Identity layer to cast it to dtype nvinfer1::ITensor* castITensor(ConversionCtx* ctx, nvinfer1::ITensor* tensor, nvinfer1::DataType dtype);