[fix]: fix bug in aten::to, when network only have aten::to layer will change input name

Signed-off-by: inocsin <[email protected]>

Author: inocsin

core/conversion/converters/converter_util.cpp

@@ -129,24 +129,24 @@ nvinfer1::ITensor* applyIdentityOp(ConversionCtx* ctx, nvinfer1::ITensor* tensor
 }
 
 nvinfer1::ITensor* castITensor(ConversionCtx* ctx, nvinfer1::ITensor* tensor, nvinfer1::DataType dtype) {
-  if (tensor->getType() != dtype) {
-    std::ostringstream tensor_id;
-    tensor_id << reinterpret_cast<int*>(tensor);
+  // No matter whether tensor->getType() == dtype, identity layer is always needed.
+  // Otherwise will change the input tensor name in aten::to converter by AssociateValueAndTensor function
+  // When the input of aten::to is network input, will cause error
+  std::ostringstream tensor_id;
+  tensor_id << reinterpret_cast<int*>(tensor);
 
-    auto id_layer = ctx->net->addIdentity(*tensor);
-    TORCHTRT_CHECK(id_layer, "Unable to create identity layer for ITensor: " << tensor_id.str());
-    auto casted_tensor = id_layer->getOutput(0);
-    casted_tensor->setType(dtype);
+  auto id_layer = ctx->net->addIdentity(*tensor);
+  TORCHTRT_CHECK(id_layer, "Unable to create identity layer for ITensor: " << tensor_id.str());
+  auto casted_tensor = id_layer->getOutput(0);
+  casted_tensor->setType(dtype);
 
-    LOG_DEBUG(ctx->logger, "Casting ITensor " << tensor_id.str() << " from " << tensor->getType() << " to " << dtype);
+  LOG_DEBUG(ctx->logger, "Casting ITensor " << tensor_id.str() << " from " << tensor->getType() << " to " << dtype);
+
+  std::stringstream ss;
+  ss << "[Cast ITensor " << tensor_id.str() << " from " << tensor->getType() << " to " << dtype << "]";
+  id_layer->setName(ss.str().c_str());
+  return casted_tensor;
 
-    std::stringstream ss;
-    ss << "[Cast ITensor " << tensor_id.str() << " from " << tensor->getType() << " to " << dtype << "]";
-    id_layer->setName(ss.str().c_str());
-    return casted_tensor;
-  } else {
-    return tensor;
-  }
 }
 
 nvinfer1::ITensor* tensor_to_const(ConversionCtx* ctx, at::Tensor t, const std::string& name) {

tests/core/conversion/converters/test_cast.cpp

@@ -135,6 +135,36 @@ TEST(Converters, ATenBoolToINT32TensorConvertsCorrectly) {
   ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6));
 }
 
+
+TEST(Converters, ATenToSingleConvertsCorrectly) {
+  const auto graph = R"IR(
+    graph(%y.1 : Tensor):
+            %4 : int = prim::Constant[value=6]()
+            %5 : bool = prim::Constant[value=0]()
+            %6 : None = prim::Constant()
+            %y0.1 : Tensor = aten::to(%y.1, %4, %5, %5, %6)
+            return (%y0.1))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+
+  torch::jit::parseIR(graph, &*g);
+
+  auto in = at::randint(1, 10, {3}, {at::kCUDA});
+
+  auto jit_in = at::clone(in);
+  auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in});
+
+  auto trt_in = at::clone(in);
+  params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in});
+
+  auto trt = trt_results[0].reshape(jit_results[0].sizes());
+  ASSERT_TRUE(jit_results[0].scalar_type() == trt.scalar_type());
+  ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6));
+}
+
+
 TEST(Converters, ATenTypeAsConvertsCorrectly) {
   const auto graph = R"IR(
       graph(%0 : Tensor,