Repository: pytorch/TensorRT

Issue:
✨[Feature] Support aten::Int.Tensor(Tensor a) -> (int)
**Is your feature request related to a problem? Please describe.**
<!--A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]-->
I think "aten::Int" is a common function. Conversion failed because converter `aten::Int.Tensor(Tensor a) -> (int)` was missing. There is a subgraph about using it.
```
const auto graph = R"IR(
      graph(%0 : Tensor):
        %1 : int = prim::Constant[value=0]()
        %1.1 : int = prim::Constant[value=1]()
        %1.2 : int = prim::Constant[value=2]()
        %2 : int = aten::size(%0, %1)
        %2.1 : int = aten::size(%0, %1.1)
        %2.2 : int = aten::size(%0, %1.2)
        %3 : Tensor = prim::NumToTensor(%2.1)
        %4 : Tensor = prim::NumToTensor(%2.2)
        %5 : Tensor = aten::mul(%3, %4)
        %6 : int = aten::Int(%5)
        %7 : int[] = prim::ListConstruct(%2, %6)
        %8 : Tensor = aten::view(%0, %7)
        return (%8))IR";
```
Input could be  at::randint(1, 10, {4, 3, 10}, {at::kCUDA})
**Describe the solution you'd like**
<!--A clear and concise description of what you want to happen.-->
Add a converter.
**Describe alternatives you've considered**
<!--A clear and concise description of any alternative solutions or features you've considered.-->

**Additional context**
<!--Add any other context or screenshots about the feature request here.-->



Discussion:
I believe this issue occurs when model tries to save static integer information for operations such as reshape. An example would be CSWin Transformer where this information is determined at model creation and stays static always, TorchScript I believe treats such info as IntTensors. 

Resolving patch:
diff --git a/noxfile.py b/noxfile.py
--- a/noxfile.py
+++ b/noxfile.py
@@ -18,7 +18,7 @@ def install_deps(session):
 
 def download_models(session, use_host_env=False):
     print("Downloading test models")
-    session.install('timm')
+    session.install("-r", os.path.join(TOP_DIR, "tests", "modules", "requirements.txt"))
     print(TOP_DIR)
     session.chdir(os.path.join(TOP_DIR, "tests", "modules"))
     if use_host_env:


Test patch:
diff --git a/tests/core/lowering/BUILD b/tests/core/lowering/BUILD
--- a/tests/core/lowering/BUILD
+++ b/tests/core/lowering/BUILD
@@ -50,6 +50,10 @@ lowering_test(
     name = "test_remove_detach_pass",
 )
 
+lowering_test(
+    name = "test_remove_unnecessary_casts",
+)
+
 lowering_test(
     name = "test_view_to_reshape_pass",
 )
@@ -81,6 +85,7 @@ test_suite(
         ":test_remove_detach_pass",
         ":test_view_to_reshape_pass",
         ":test_remove_dropout_pass",
+        ":test_remove_unnecessary_casts",
         ":test_reduce_to_pass",
         ":test_reduce_gelu",
         ":test_unpack_hardswish",
diff --git a/tests/core/lowering/test_remove_unnecessary_casts.cpp b/tests/core/lowering/test_remove_unnecessary_casts.cpp
new file mode 100644
--- /dev/null
+++ b/tests/core/lowering/test_remove_unnecessary_casts.cpp
@@ -0,0 +1,155 @@
+#include <string>
+#include "core/compiler.h"
+#include "core/lowering/passes/passes.h"
+#include "gtest/gtest.h"
+#include "tests/util/util.h"
+#include "torch/csrc/jit/ir/irparser.h"
+#include "torch/csrc/jit/ir/subgraph_matcher.h"
+
+TEST(LoweringPasses, RemoveUnnecessaryCastIntCorrectly) {
+  std::string source_graph = R"IR(
+    graph(%1: int):
+      %2: Tensor = aten::NumToTensor(%1)
+      %3: int = aten::Int(%2)
+      %4: int = aten::add(%3, %3, %3)
+      return (%4))IR";
+  std::string target_graph = R"IR(
+    graph(%1: int):
+      %4: int = aten::add(%1, %1, %1)
+      return (%4))IR";
+
+  torch_tensorrt::core::util::logging::get_logger().set_reportable_log_level(
+      torch_tensorrt::core::util::logging::LogLevel::kGRAPH);
+  auto sg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(source_graph, sg.get());
+  torch_tensorrt::core::lowering::passes::RemoveUnnecessaryCasts(sg);
+
+  auto tg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(target_graph, tg.get());
+
+  ASSERT_TRUE(!torch::jit::findPatternMatches(*tg, *sg).empty());
+}
+
+TEST(LoweringPasses, RemoveUnnecessaryCastFloatCorrectly) {
+  std::string source_graph = R"IR(
+    graph(%1: float):
+      %2: Tensor = aten::NumToTensor(%1)
+      %3: float = aten::Float(%2)
+      %4: float = aten::add(%3, %3, %3)
+      return (%3))IR";
+  std::string target_graph = R"IR(
+    graph(%1: float):
+      %4: float = aten::add(%1, %1, %1)
+      return (%4))IR";
+
+  torch_tensorrt::core::util::logging::get_logger().set_reportable_log_level(
+      torch_tensorrt::core::util::logging::LogLevel::kGRAPH);
+  auto sg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(source_graph, sg.get());
+  torch_tensorrt::core::lowering::passes::RemoveUnnecessaryCasts(sg);
+
+  auto tg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(target_graph, tg.get());
+
+  ASSERT_TRUE(!torch::jit::findPatternMatches(*tg, *sg).empty());
+}
+
+TEST(LoweringPasses, RemoveUnnecessaryCastBoolCorrectly) {
+  std::string source_graph = R"IR(
+    graph(%1: bool):
+      %2: Tensor = aten::NumToTensor(%1)
+      %3: bool = aten::Bool(%2)
+      %4: bool = aten::__and__(%3, %3)
+      return (%3))IR";
+  std::string target_graph = R"IR(
+    graph(%1: bool):
+      %4: bool = aten::__and__(%1, %1)
+      return (%4))IR";
+
+  torch_tensorrt::core::util::logging::get_logger().set_reportable_log_level(
+      torch_tensorrt::core::util::logging::LogLevel::kGRAPH);
+  auto sg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(source_graph, sg.get());
+  torch_tensorrt::core::lowering::passes::RemoveUnnecessaryCasts(sg);
+
+  auto tg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(target_graph, tg.get());
+
+  ASSERT_TRUE(!torch::jit::findPatternMatches(*tg, *sg).empty());
+}
+
+TEST(LoweringPasses, RemoveSingleUse0DTensorsIntCorrectly) {
+  std::string source_graph = R"IR(
+    graph(%0: int):
+      %1: Tensor = prim::Constant[value=[8]]()
+      %2: int = prim::Constant[value=1]()
+      %3: Tensor = prim::NumToTensor(%0)
+      %4: Tensor = aten::add(%1, %3, %2)
+      %5: int = aten::Int(%4)
+      %6: int = aten::add(%5, %5)
+      return (%6))IR";
+  std::string target_graph = R"IR(
+    graph(%0: int):
+      %1: int = prim::Constant[value=8]()
+      %4: int = aten::add(%1, %0)
+      %6: int = aten::add(%4, %4)
+      return (%6))IR";
+
+  torch_tensorrt::core::util::logging::get_logger().set_reportable_log_level(
+      torch_tensorrt::core::util::logging::LogLevel::kGRAPH);
+  auto sg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(source_graph, sg.get());
+
+  auto first_op = *(sg->block()->nodes().begin());
+  torch::jit::WithInsertPoint guard(first_op);
+  torch::jit::Value* r = sg->insertConstant(c10::scalar_to_tensor(8), c10::nullopt, first_op->scope());
+  r->copyMetadata(first_op->output());
+  r->setType(c10::TensorType::get());
+  first_op->output()->replaceAllUsesWith(r);
+  first_op->destroy();
+
+  torch_tensorrt::core::lowering::passes::RemoveSingleUse0DTensors(sg);
+
+  auto tg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(target_graph, tg.get());
+
+  ASSERT_TRUE(!torch::jit::findPatternMatches(*tg, *sg).empty());
+}
+
+TEST(LoweringPasses, RemoveSingleUse0DTensorsFloatCorrectly) {
+  std::string source_graph = R"IR(
+    graph(%0: float):
+      %1: Tensor = prim::Constant[value=[8.]]()
+      %2: float = prim::Constant[value=1.]()
+      %3: Tensor = prim::NumToTensor(%0)
+      %4: Tensor = aten::add(%1, %3, %2)
+      %5: float = aten::Float(%4)
+      %6: float = aten::add(%5, %5)
+      return (%6))IR";
+  std::string target_graph = R"IR(
+    graph(%0: float):
+      %1: float = prim::Constant[value=8.]()
+      %4: float = aten::add(%1, %0)
+      %6: float = aten::add(%4, %4)
+      return (%6))IR";
+
+  torch_tensorrt::core::util::logging::get_logger().set_reportable_log_level(
+      torch_tensorrt::core::util::logging::LogLevel::kGRAPH);
+  auto sg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(source_graph, sg.get());
+
+  auto first_op = *(sg->block()->nodes().begin());
+  torch::jit::WithInsertPoint guard(first_op);
+  torch::jit::Value* r = sg->insertConstant(c10::scalar_to_tensor(8.0), c10::nullopt, first_op->scope());
+  r->copyMetadata(first_op->output());
+  r->setType(c10::TensorType::get());
+  first_op->output()->replaceAllUsesWith(r);
+  first_op->destroy();
+
+  torch_tensorrt::core::lowering::passes::RemoveSingleUse0DTensors(sg);
+
+  auto tg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(target_graph, tg.get());
+
+  ASSERT_TRUE(!torch::jit::findPatternMatches(*tg, *sg).empty());
+}
\ No newline at end of file
diff --git a/tests/cpp/cpp_api_test.h b/tests/cpp/cpp_api_test.h
--- a/tests/cpp/cpp_api_test.h
+++ b/tests/cpp/cpp_api_test.h
@@ -6,12 +6,12 @@
 #include "torch/script.h"
 #include "torch_tensorrt/torch_tensorrt.h"
 
-using PathAndInSize = std::tuple<std::string, std::vector<std::vector<int64_t>>, float>;
+using PathAndInput = std::tuple<std::string, std::vector<std::vector<int64_t>>, std::vector<c10::ScalarType>, float>;
 
-class CppAPITests : public testing::TestWithParam<PathAndInSize> {
+class CppAPITests : public testing::TestWithParam<PathAndInput> {
  public:
   void SetUp() override {
-    PathAndInSize params = GetParam();
+    PathAndInput params = GetParam();
     std::string path = std::get<0>(params);
     try {
       // Deserialize the ScriptModule from a file using torch::jit::load().
@@ -21,7 +21,8 @@ class CppAPITests : public testing::TestWithParam<PathAndInSize> {
       ASSERT_TRUE(false);
     }
     input_shapes = std::get<1>(params);
-    threshold = std::get<2>(params);
+    input_types = std::get<2>(params);
+    threshold = std::get<3>(params);
   }
 
   void TearDown() {
@@ -32,5 +33,6 @@ class CppAPITests : public testing::TestWithParam<PathAndInSize> {
  protected:
   torch::jit::script::Module mod;
   std::vector<std::vector<int64_t>> input_shapes;
+  std::vector<c10::ScalarType> input_types;
   float threshold;
 };
diff --git a/tests/cpp/test_compiled_modules.cpp b/tests/cpp/test_compiled_modules.cpp
--- a/tests/cpp/test_compiled_modules.cpp
+++ b/tests/cpp/test_compiled_modules.cpp
@@ -3,20 +3,42 @@
 TEST_P(CppAPITests, CompiledModuleIsClose) {
   std::vector<torch::jit::IValue> jit_inputs_ivalues;
   std::vector<torch::jit::IValue> trt_inputs_ivalues;
-  for (auto in_shape : input_shapes) {
-    auto in = at::randint(5, in_shape, {at::kCUDA});
+  std::vector<torch_tensorrt::Input> shapes;
+  for (uint64_t i = 0; i < input_shapes.size(); i++) {
+    auto in = at::randint(5, input_shapes[i], {at::kCUDA}).to(input_types[i]);
     jit_inputs_ivalues.push_back(in.clone());
     trt_inputs_ivalues.push_back(in.clone());
+    auto in_spec = torch_tensorrt::Input(input_shapes[i]);
+    in_spec.dtype = input_types[i];
+    shapes.push_back(in_spec);
+    std::cout << in_spec << std::endl;
   }
 
   torch::jit::IValue jit_results_ivalues = torch_tensorrt::tests::util::RunModuleForward(mod, jit_inputs_ivalues);
   std::vector<at::Tensor> jit_results;
-  jit_results.push_back(jit_results_ivalues.toTensor());
+  if (jit_results_ivalues.isTuple()) {
+    auto tuple = jit_results_ivalues.toTuple();
+    for (auto t : tuple->elements()) {
+      jit_results.push_back(t.toTensor());
+    }
+  } else {
+    jit_results.push_back(jit_results_ivalues.toTensor());
+  }
+
+  auto spec = torch_tensorrt::ts::CompileSpec(shapes);
+  spec.truncate_long_and_double = true;
 
-  auto trt_mod = torch_tensorrt::ts::compile(mod, input_shapes);
+  auto trt_mod = torch_tensorrt::ts::compile(mod, spec);
   torch::jit::IValue trt_results_ivalues = torch_tensorrt::tests::util::RunModuleForward(trt_mod, trt_inputs_ivalues);
   std::vector<at::Tensor> trt_results;
-  trt_results.push_back(trt_results_ivalues.toTensor());
+  if (trt_results_ivalues.isTuple()) {
+    auto tuple = trt_results_ivalues.toTuple();
+    for (auto t : tuple->elements()) {
+      trt_results.push_back(t.toTensor());
+    }
+  } else {
+    trt_results.push_back(trt_results_ivalues.toTensor());
+  }
 
   for (size_t i = 0; i < trt_results.size(); i++) {
     ASSERT_TRUE(
@@ -30,13 +52,14 @@ INSTANTIATE_TEST_SUITE_P(
     CompiledModuleForwardIsCloseSuite,
     CppAPITests,
     testing::Values(
-        PathAndInSize({"tests/modules/resnet18_traced.jit.pt", {{1, 3, 224, 224}}, 2e-5}),
-        PathAndInSize({"tests/modules/resnet50_traced.jit.pt", {{1, 3, 224, 224}}, 2e-5}),
-        PathAndInSize({"tests/modules/mobilenet_v2_traced.jit.pt", {{1, 3, 224, 224}}, 2e-5}),
-        PathAndInSize({"tests/modules/resnet18_scripted.jit.pt", {{1, 3, 224, 224}}, 2e-5}),
-        PathAndInSize({"tests/modules/resnet50_scripted.jit.pt", {{1, 3, 224, 224}}, 2e-5}),
-        PathAndInSize({"tests/modules/mobilenet_v2_scripted.jit.pt", {{1, 3, 224, 224}}, 2e-5}),
-        PathAndInSize({"tests/modules/efficientnet_b0_scripted.jit.pt", {{1, 3, 224, 224}}, 8e-3}),
-        PathAndInSize({"tests/modules/vit_scripted.jit.pt", {{1, 3, 224, 224}}, 8e-2})));
+        PathAndInput({"tests/modules/resnet18_traced.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 2e-5}),
+        PathAndInput({"tests/modules/resnet50_traced.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 2e-5}),
+        PathAndInput({"tests/modules/mobilenet_v2_traced.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 2e-5}),
+        PathAndInput({"tests/modules/resnet18_scripted.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 2e-5}),
+        PathAndInput({"tests/modules/resnet50_scripted.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 2e-5}),
+        PathAndInput({"tests/modules/mobilenet_v2_scripted.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 2e-5}),
+        PathAndInput({"tests/modules/efficientnet_b0_scripted.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 8e-3}),
+        PathAndInput({"tests/modules/bert_base_uncased_traced.jit.pt", {{1, 14}, {1, 14}}, {at::kInt, at::kInt}, 8e-2}),
+        PathAndInput({"tests/modules/vit_scripted.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 8e-2})));
 
 #endif
diff --git a/tests/cpp/test_default_input_types.cpp b/tests/cpp/test_default_input_types.cpp
--- a/tests/cpp/test_default_input_types.cpp
+++ b/tests/cpp/test_default_input_types.cpp
@@ -78,7 +78,7 @@ TEST_P(CppAPITests, InputsRespectUserSettingFP16WeightsFP32In) {
   }
 
   auto in = torch_tensorrt::Input(input_shapes[0]);
-  in.dtype = torch::kF32;
+  in.dtype = torch::kFloat;
   auto spec = torch_tensorrt::ts::CompileSpec({in});
   spec.enabled_precisions.insert(torch_tensorrt::DataType::kHalf);
 
@@ -116,4 +116,5 @@ TEST_P(CppAPITests, InputsRespectUserSettingFP32WeightsFP16In) {
 INSTANTIATE_TEST_SUITE_P(
     CompiledModuleForwardIsCloseSuite,
     CppAPITests,
-    testing::Values(PathAndInSize({"tests/modules/resnet18_traced.jit.pt", {{1, 3, 224, 224}}, 2e-5})));
+    testing::Values(
+        PathAndInput({"tests/modules/resnet18_traced.jit.pt", {{1, 3, 224, 224}}, {at::kFloat} /*unused*/, 2e-5})));
diff --git a/tests/cpp/test_example_tensors.cpp b/tests/cpp/test_example_tensors.cpp
--- a/tests/cpp/test_example_tensors.cpp
+++ b/tests/cpp/test_example_tensors.cpp
@@ -3,8 +3,8 @@
 TEST_P(CppAPITests, InputsFromTensors) {
   std::vector<torch::jit::IValue> jit_inputs_ivalues;
   std::vector<torch::jit::IValue> trt_inputs_ivalues;
-  for (auto in_shape : input_shapes) {
-    auto in = at::randn(in_shape, {at::kCUDA});
+  for (uint64_t i = 0; i < input_shapes.size(); i++) {
+    auto in = at::randn(input_shapes[i], {at::kCUDA}).to(input_types[i]);
     jit_inputs_ivalues.push_back(in.clone());
     trt_inputs_ivalues.push_back(in.clone());
   }
@@ -20,4 +20,4 @@ TEST_P(CppAPITests, InputsFromTensors) {
 INSTANTIATE_TEST_SUITE_P(
     CompiledModuleForwardIsCloseSuite,
     CppAPITests,
-    testing::Values(PathAndInSize({"tests/modules/resnet18_traced.jit.pt", {{1, 3, 224, 224}}, 2e-5})));
+    testing::Values(PathAndInput({"tests/modules/resnet18_traced.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 2e-5})));
diff --git a/tests/cpp/test_modules_as_engines.cpp b/tests/cpp/test_modules_as_engines.cpp
--- a/tests/cpp/test_modules_as_engines.cpp
+++ b/tests/cpp/test_modules_as_engines.cpp
@@ -4,8 +4,8 @@
 TEST_P(CppAPITests, ModuleAsEngineIsClose) {
   std::vector<at::Tensor> inputs;
   std::vector<torch::jit::IValue> inputs_ivalues;
-  for (auto in_shape : input_shapes) {
-    inputs.push_back(at::randint(5, in_shape, {at::kCUDA}));
+  for (uint64_t i = 0; i < input_shapes.size(); i++) {
+    inputs.push_back(at::randint(5, input_shapes[i], {at::kCUDA}).to(input_types[i]));
     inputs_ivalues.push_back(inputs[inputs.size() - 1].clone());
   }
 
@@ -21,8 +21,8 @@ TEST_P(CppAPITests, ModuleAsEngineIsClose) {
 TEST_P(CppAPITests, ModuleToEngineToModuleIsClose) {
   std::vector<at::Tensor> inputs;
   std::vector<torch::jit::IValue> inputs_ivalues;
-  for (auto in_shape : input_shapes) {
-    inputs.push_back(at::randint(5, in_shape, {at::kCUDA}));
+  for (uint64_t i = 0; i < input_shapes.size(); i++) {
+    inputs.push_back(at::randint(5, input_shapes[i], {at::kCUDA}).to(input_types[i]));
     inputs_ivalues.push_back(inputs[inputs.size() - 1].clone());
   }
 
@@ -57,13 +57,13 @@ INSTANTIATE_TEST_SUITE_P(
     ModuleAsEngineForwardIsCloseSuite,
     CppAPITests,
     testing::Values(
-        PathAndInSize({"tests/modules/resnet18_traced.jit.pt", {{1, 3, 224, 224}}, 2e-5}),
-        PathAndInSize({"tests/modules/resnet50_traced.jit.pt", {{1, 3, 224, 224}}, 2e-5}),
-        PathAndInSize({"tests/modules/mobilenet_v2_traced.jit.pt", {{1, 3, 224, 224}}, 2e-5}),
-        PathAndInSize({"tests/modules/resnet18_scripted.jit.pt", {{1, 3, 224, 224}}, 2e-5}),
-        PathAndInSize({"tests/modules/resnet50_scripted.jit.pt", {{1, 3, 224, 224}}, 2e-5}),
-        PathAndInSize({"tests/modules/mobilenet_v2_scripted.jit.pt", {{1, 3, 224, 224}}, 2e-5}),
-        PathAndInSize({"tests/modules/efficientnet_b0_scripted.jit.pt", {{1, 3, 224, 224}}, 2e-5}),
-        PathAndInSize({"tests/modules/vit_scripted.jit.pt", {{1, 3, 224, 224}}, 8e-2})));
+        PathAndInput({"tests/modules/resnet18_traced.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 2e-5}),
+        PathAndInput({"tests/modules/resnet50_traced.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 2e-5}),
+        PathAndInput({"tests/modules/mobilenet_v2_traced.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 2e-5}),
+        PathAndInput({"tests/modules/resnet18_scripted.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 2e-5}),
+        PathAndInput({"tests/modules/resnet50_scripted.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 2e-5}),
+        PathAndInput({"tests/modules/mobilenet_v2_scripted.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 2e-5}),
+        PathAndInput({"tests/modules/efficientnet_b0_scripted.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 2e-5}),
+        PathAndInput({"tests/modules/vit_scripted.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 8e-2})));
 
 #endif
diff --git a/tests/cpp/test_multi_gpu_serde.cpp b/tests/cpp/test_multi_gpu_serde.cpp
--- a/tests/cpp/test_multi_gpu_serde.cpp
+++ b/tests/cpp/test_multi_gpu_serde.cpp
@@ -4,8 +4,8 @@
 TEST_P(CppAPITests, CompiledModuleIsClose) {
   std::vector<torch::jit::IValue> jit_inputs_ivalues;
   std::vector<torch::jit::IValue> trt_inputs_ivalues;
-  for (auto in_shape : input_shapes) {
-    auto in = at::randint(5, in_shape, {at::kCUDA});
+  for (uint64_t i = 0; i < input_shapes.size(); i++) {
+    auto in = at::randint(5, input_shapes[i], {at::kCUDA}).to(input_types[i]);
     jit_inputs_ivalues.push_back(in.clone());
     trt_inputs_ivalues.push_back(in.clone());
   }
@@ -31,4 +31,4 @@ TEST_P(CppAPITests, CompiledModuleIsClose) {
 INSTANTIATE_TEST_SUITE_P(
     CompiledModuleForwardIsCloseSuite,
     CppAPITests,
-    testing::Values(PathAndInSize({"tests/modules/resnet18_traced.jit.pt", {{1, 3, 224, 224}}, 2e-5})));
\ No newline at end of file
+    testing::Values(PathAndInput({"tests/modules/resnet18_traced.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 2e-5})));
\ No newline at end of file
diff --git a/tests/cpp/test_serialization.cpp b/tests/cpp/test_serialization.cpp
--- a/tests/cpp/test_serialization.cpp
+++ b/tests/cpp/test_serialization.cpp
@@ -21,8 +21,8 @@ std::vector<torch_tensorrt::Input> toInputRangesDynamic(std::vector<std::vector<
 TEST_P(CppAPITests, SerializedModuleIsStillCorrect) {
   std::vector<torch::jit::IValue> post_serialized_inputs_ivalues;
   std::vector<torch::jit::IValue> pre_serialized_inputs_ivalues;
-  for (auto in_shape : input_shapes) {
-    auto in = at::randint(5, in_shape, {at::kCUDA});
+  for (uint64_t i = 0; i < input_shapes.size(); i++) {
+    auto in = at::randint(5, input_shapes[i], {at::kCUDA}).to(input_types[i]);
     post_serialized_inputs_ivalues.push_back(in.clone());
     pre_serialized_inputs_ivalues.push_back(in.clone());
   }
@@ -50,8 +50,8 @@ TEST_P(CppAPITests, SerializedModuleIsStillCorrect) {
 TEST_P(CppAPITests, SerializedDynamicModuleIsStillCorrect) {
   std::vector<torch::jit::IValue> post_serialized_inputs_ivalues;
   std::vector<torch::jit::IValue> pre_serialized_inputs_ivalues;
-  for (auto in_shape : input_shapes) {
-    auto in = at::randint(5, in_shape, {at::kCUDA});
+  for (uint64_t i = 0; i < input_shapes.size(); i++) {
+    auto in = at::randint(5, input_shapes[i], {at::kCUDA}).to(input_types[i]);
     post_serialized_inputs_ivalues.push_back(in.clone());
     pre_serialized_inputs_ivalues.push_back(in.clone());
   }
@@ -81,5 +81,5 @@ INSTANTIATE_TEST_SUITE_P(
     CompiledModuleForwardIsCloseSuite,
     CppAPITests,
     testing::Values(
-        PathAndInSize({"tests/modules/resnet18_traced.jit.pt", {{1, 3, 224, 224}}, 2e-5}),
-        PathAndInSize({"tests/modules/pooling_traced.jit.pt", {{1, 3, 10, 10}}, 2e-5})));
+        PathAndInput({"tests/modules/resnet18_traced.jit.pt", {{1, 3, 224, 224}}, {at::kFloat}, 2e-5}),
+        PathAndInput({"tests/modules/pooling_traced.jit.pt", {{1, 3, 10, 10}}, {at::kFloat}, 2e-5})));
diff --git a/tests/modules/hub.py b/tests/modules/hub.py
--- a/tests/modules/hub.py
+++ b/tests/modules/hub.py
@@ -3,6 +3,7 @@
 import torch.nn.functional as F
 import torchvision.models as models
 import timm
+from transformers import BertModel, BertTokenizer, BertConfig
 
 torch.hub._validate_not_a_forked_repo = lambda a, b, c: True
 
@@ -189,3 +190,30 @@ def forward(self, x):
 conditional_model = FallbackIf().eval().cuda()
 conditional_script_model = torch.jit.script(conditional_model)
 torch.jit.save(conditional_script_model, "conditional_scripted.jit.pt")
+
+enc = BertTokenizer.from_pretrained("bert-base-uncased")
+text = "[CLS] Who was Jim Henson ? [SEP] Jim Henson was a puppeteer [SEP]"
+tokenized_text = enc.tokenize(text)
+masked_index = 8
+tokenized_text[masked_index] = "[MASK]"
+indexed_tokens = enc.convert_tokens_to_ids(tokenized_text)
+segments_ids = [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1]
+tokens_tensor = torch.tensor([indexed_tokens])
+segments_tensors = torch.tensor([segments_ids])
+dummy_input = [tokens_tensor, segments_tensors]
+
+config = BertConfig(
+    vocab_size_or_config_json_file=32000,
+    hidden_size=768,
+    num_hidden_layers=12,
+    num_attention_heads=12,
+    intermediate_size=3072,
+    torchscript=True,
+)
+
+model = BertModel(config)
+model.eval()
+model = BertModel.from_pretrained("bert-base-uncased", torchscript=True)
+
+traced_model = torch.jit.trace(model, [tokens_tensor, segments_tensors])
+torch.jit.save(traced_model, "bert_base_uncased_traced.jit.pt")
diff --git a/tests/modules/requirements.txt b/tests/modules/requirements.txt
--- a/tests/modules/requirements.txt
+++ b/tests/modules/requirements.txt
@@ -1,3 +1,4 @@
 -f https://download.pytorch.org/whl/torch_stable.html
+#torch==1.10.0+cu113
 timm==v0.4.12
-torch==1.10.0+cu113
+transformers==4.17.0
\ No newline at end of file