Merge pull request #395 from guoruoqian/unsqueeze_fix_bug

fix a bug when unsqueeze(self, dim)'s dim is negative, the result of …

Author: narendasan

core/conversion/converters/impl/unsqueeze.cpp

@@ -21,8 +21,13 @@ auto unsqueeze_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns().p
        auto dim = args[1].unwrapToInt();
 
        auto selfDim = util::toVec(self->getDimensions());
+       int64_t nbDims = selfDim.size();
+       TRTORCH_CHECK(
+           dim <= nbDims && dim >= -(nbDims + 1),
+           "Dimension out of range (expected to be in range of [" << -(nbDims + 1) << ", " << nbDims << "], but got "
+                                                                  << dim << ")");
        if (dim < 0) {
-         dim = selfDim.size() + dim;
+         dim = nbDims + 1 + dim;
        }
 
        auto shuffle_layer = ctx->net->addShuffle(*self);

tests/core/conversion/converters/test_unsqueeze.cpp

@@ -22,5 +22,26 @@ TEST(Converters, ATenUnsqueezeConvertsCorrectly) {
   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, ATenUnsqueezeNegativeDimConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %1 : int = prim::Constant[value=-4]()
+        %2 : Tensor = aten::unsqueeze(%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));
 }