fix: fix the renaming error in squeeze converter

Signed-off-by: Bo Wang <[email protected]>

Author: bowang007

core/conversion/converters/impl/squeeze.cpp

@@ -26,7 +26,8 @@ auto squeeze_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pa
        }
 
        if (selfDim[dim] != 1) {
-         auto out = ctx->AssociateValueAndTensor(n->outputs()[0], self);
+         auto id_layer = ctx->net->addIdentity(*self);
+         auto out = ctx->AssociateValueAndTensor(n->outputs()[0], id_layer->getOutput(0));
 
          LOG_DEBUG("Output tensor shape: " << out->getDimensions());
 

tests/core/conversion/converters/test_squeeze.cpp

@@ -53,6 +53,32 @@ TEST(Converters, ATenSqueezeDontNeedSqueezeConvertsCorrectly) {
   params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
   auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in, trt_in_add});
 
+  ASSERT_TRUE(
+      torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
+}
+
+TEST(Converters, ATenSqueezeNeedIdentityConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %2 : int = prim::Constant[value=1]()
+        %3 : Tensor = aten::squeeze(%0, %2)
+        return (%3))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, &*g);
+
+  auto in = at::randint(1, 10, {2, 3, 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(jit_in);
+
+  params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in});
+
   ASSERT_TRUE(
       torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
 }