chore: Add testcase

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

core/conversion/converters/impl/shuffle.cpp

@@ -78,7 +78,7 @@ static auto shuffle_registrations TORCHTRT_UNUSED =
                  concat_layer->setAxis(static_cast<int32_t>(0));
                  shape_tensor = concat_layer->getOutput(0);
                } else {
-                 auto new_shape = torch::reshape(torch::rand(in_shape), args[1].unwrapToIntList().vec()).sizes().vec();
+                 new_shape = torch::reshape(torch::rand(in_shape), args[1].unwrapToIntList().vec()).sizes().vec();
                }
                auto shuffle = ctx->net->addShuffle(*in);
                shuffle->setName(util::node_info(n).c_str());

core/conversion/evaluators/prim.cpp

@@ -1,12 +1,11 @@
 #include <limits>
 
-#include "torch/csrc/jit/ir/ir.h"
-//#include "torch/csrc/jit/ir/constants.h"
 #include "ATen/core/List.h"
 #include "ATen/core/functional.h"
 #include "ATen/core/ivalue.h"
 #include "ATen/core/stack.h"
 #include "c10/util/intrusive_ptr.h"
+#include "torch/csrc/jit/ir/ir.h"
 #include "torch/torch.h"
 
 #include "core/conversion/evaluators/eval_macros.h"
@@ -111,8 +110,20 @@ auto prim_registrations =
                        tensor_holder.hold_tensor(itensor);
                        auto ival = c10::IValue(std::move(c10::make_intrusive<TensorContainer>(tensor_holder)));
                        list.emplace_back(std::move(ival));
+                     } else if (args.at(in).IValue()->isDouble()) {
+                       auto itensor = torch_tensorrt::core::conversion::converters::tensor_to_const(
+                           ctx, torch::tensor({args.at(in).unwrapToDouble()}).to(torch::kFloat));
+                       auto tensor_holder = TensorContainer();
+                       tensor_holder.hold_tensor(itensor);
+                       auto ival = c10::IValue(std::move(c10::make_intrusive<TensorContainer>(tensor_holder)));
+                       list.emplace_back(std::move(ival));
                      } else {
-                       list.emplace_back(std::move(args.at(in).unwrapToTensor()));
+                       auto itensor = torch_tensorrt::core::conversion::converters::tensor_to_const(
+                           ctx, std::move(args.at(in).unwrapToTensor()));
+                       auto tensor_holder = TensorContainer();
+                       tensor_holder.hold_tensor(itensor);
+                       auto ival = c10::IValue(std::move(c10::make_intrusive<TensorContainer>(tensor_holder)));
+                       list.emplace_back(std::move(ival));
                      }
                    }
                  }

py/requirements.txt

@@ -1,7 +1,6 @@
 numpy
 pybind11==2.6.2
---extra-index-url https://download.pytorch.org/whl/nightly/cu117
-torch==2.0.0.dev20230103+cu117
-torchvision==0.15.0.dev20230103+cu117
+torch==1.13.0
+torchvision==0.14.0
 --extra-index-url https://pypi.ngc.nvidia.com
 tensorrt==8.5.1.7

tests/cpp/BUILD

@@ -16,6 +16,7 @@ test_suite(
         ":test_compiled_modules",
         ":test_default_input_types",
         ":test_dynamic_fallback",
+        ":test_dynamic_size",
         ":test_example_tensors",
         ":test_module_fallback",
         ":test_modules_as_engines",
@@ -32,6 +33,7 @@ test_suite(
         ":test_compiled_modules",
         ":test_default_input_types",
         ":test_dynamic_fallback",
+        ":test_dynamic_size",
         ":test_example_tensors",
         ":test_module_fallback",
         ":test_modules_as_engines",
@@ -142,6 +144,18 @@ cc_test(
     }),
 )
 
+cc_test(
+    name = "test_dynamic_size",
+    srcs = ["test_dynamic_size.cpp"],
+    deps = [
+        "//tests/util",
+        "@googletest//:gtest_main",
+    ] + select({
+        ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
+        "//conditions:default": ["@libtorch//:libtorch"],
+    }),
+)
+
 cc_test(
     name = "test_collections",
     srcs = ["test_collections.cpp"],

tests/cpp/test_dynamic_size.cpp

@@ -0,0 +1,35 @@
+#include <torch/torch.h>
+#include <string>
+#include "core/compiler.h"
+#include "gtest/gtest.h"
+#include "tests/util/util.h"
+#include "torch/csrc/jit/ir/irparser.h"
+
+TEST(Converters, ATenResizeDynamicInputCorrectly) {
+  const auto graph = R"IR(
+    graph(%x : Tensor):
+          %3 : int = prim::Constant[value=0]()
+          %2 : int = prim::Constant[value=-1]()
+          %28 : int = aten::size(%x, %3)
+          %30 : int[] = prim::ListConstruct(%28, %2)
+          %6 : Tensor = aten::reshape(%x, %30)
+          return (%6))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+
+  torch::jit::parseIR(graph, g.get());
+
+  auto in = at::randint(1, 10, {16, 3, 2}, {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::RunGraphEngineDynamic(g, params, {in}, true);
+
+  auto trt = trt_results[0].reshape(jit_results[0].sizes());
+
+  ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6));
+}