chore: Fix missing return casted tensor

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>
pytorch · Jul 7, 2022 · 5da47c1 · 5da47c1
1 parent 7266cad
commit 5da47c1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/core/conversion/converters/converter_util.cpp b/core/conversion/converters/converter_util.cpp
@@ -137,7 +137,7 @@ nvinfer1::ITensor* castITensor(ConversionCtx* ctx, nvinfer1::ITensor* tensor, nv
     TORCHTRT_CHECK(id_layer, "Unable to create identity layer for ITensor: " << tensor_id.str());
     // layer->setOutputType should be used for casting and not manually setting output_tensor->setType()
     id_layer->setOutputType(0, dtype);
-
+    auto casted_tensor = id_layer->getOutput(0);
     LOG_DEBUG(ctx->logger, "Casting ITensor " << tensor_id.str() << " from " << tensor->getType() << " to " << dtype);
 
     std::stringstream ss;