some tensor like at::randint(1,10,{2,3,3}) dont need to squeeze, add a judgment condition to handle this situation

Signed-off-by: Ruoqian Guo <[email protected]>

Author: ruoqianguo

core/conversion/converters/impl/squeeze.cpp

@@ -25,6 +25,15 @@ auto squeeze_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns().pat
          dim = selfDim.size() + dim;
        }
 
+       if (selfDim[dim] != 1) {
+         auto identity = ctx->net->addIdentity(*self);
+         auto out = ctx->AssociateValueAndTensor(n->outputs()[0], identity->getOutput(0));
+
+         LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+
+         return true;
+       }
+
        auto shuffle_layer = ctx->net->addShuffle(*self);
        TRTORCH_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
        shuffle_layer->setReshapeDimensions(util::squeezeDims(self->getDimensions(), dim));

tests/core/conversion/converters/test_squeeze.cpp

@@ -22,5 +22,26 @@ TEST(Converters, ATenSqueezeConvertsCorrectly) {
   params = trtorch::core::conversion::get_named_params(g->inputs(), {});
   auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in});
 
+  ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
+}
+
+TEST(Converters, ATenSqueezeDontNeedSqueezeConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %1 : int = prim::Constant[value=1]()
+        %2 : Tensor = aten::squeeze(%0, %1)
+        return (%2))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 params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto jit_results = trtorch::tests::util::RunGraph(g, params, {in});
+
+  params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in});
+
   ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
 }