refactor: way simpler module test

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

tests/cpp/test_module_fallback.cpp

@@ -4,7 +4,7 @@
 #include "torch/script.h"
 #include "trtorch/trtorch.h"
 
-TEST(CppAPITests, LowerResNetModuleFallbackCorrectly) {
+TEST(CppAPITest, ResNetModuleFallbacksCorrectly) {
   torch::jit::script::Module mod;
   try {
     mod = torch::jit::load("tests/modules/resnet18_traced.jit.pt");
@@ -22,9 +22,7 @@ TEST(CppAPITests, LowerResNetModuleFallbackCorrectly) {
     trt_inputs_ivalues.push_back(in.clone());
   }
 
-  std::vector<trtorch::CompileSpec::Input> input_ranges{
-      trtorch::CompileSpec::Input(std::vector<int64_t>({1, 3, 224, 224}))};
-  trtorch::CompileSpec cfg(input_ranges);
+  trtorch::CompileSpec cfg(input_shapes);
   cfg.torch_fallback.enabled = true;
   cfg.torch_fallback.forced_fallback_modules.push_back("torchvision.models.resnet.BasicBlock");
 
@@ -34,7 +32,7 @@ TEST(CppAPITests, LowerResNetModuleFallbackCorrectly) {
   ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results, trt_results, 2e-6));
 }
 
-TEST(CppAPITests, LowerAndPartitionMobileNetModuleFallbackCorrectly) {
+TEST(CppAPITest, MobileNetModuleFallbacksCorrectlyWithOneEngine) {
   torch::jit::script::Module mod;
   try {
     mod = torch::jit::load("tests/modules/mobilenet_v2_traced.jit.pt");
@@ -52,16 +50,24 @@ TEST(CppAPITests, LowerAndPartitionMobileNetModuleFallbackCorrectly) {
     trt_inputs_ivalues.push_back(in.clone());
   }
 
-  std::vector<trtorch::CompileSpec::Input> input_ranges{
-      trtorch::CompileSpec::Input(std::vector<int64_t>({1, 3, 224, 224}))};
-  trtorch::CompileSpec cfg(input_ranges);
+  trtorch::CompileSpec cfg(input_shapes);
   cfg.torch_fallback.enabled = true;
   cfg.torch_fallback.min_block_size = 5;
   cfg.torch_fallback.forced_fallback_modules.push_back("torchvision.models.mobilenetv2.ConvBNActivation");
 
   auto jit_results = mod.forward(jit_inputs_ivalues).toTensor();
   auto trt_mod = trtorch::CompileGraph(mod, cfg);
 
+  auto g = trt_mod.get_method("forward").graph();
+  auto nodes = g->block()->nodes();
+  std::size_t trt_count = 0;
+  for (const auto n : nodes) {
+    if (n->kind().toQualString() == std::string("tensorrt::execute_engine")) {
+      trt_count++;
+    }
+  }
+  ASSERT_TRUE(trt_count == 1);
+
   auto trt_results = trt_mod.forward(trt_inputs_ivalues).toTensor();
   ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results, trt_results, 2e-6));
 }