Merge pull request #657 from NVIDIA/fix_tests

fix: Fix modules_as_engines test case to use trt_mod instead of pyt_mod

Author: narendasan

tests/cpp/test_module_fallback.cpp

@@ -7,7 +7,7 @@
 TEST(CppAPITest, ResNetModuleFallbacksCorrectly) {
   torch::jit::script::Module mod;
   try {
-    mod = torch::jit::load("tests/modules/resnet18_traced.jit.pt");
+    mod = torch::jit::load("tests/modules/resnet18_scripted.jit.pt");
   } catch (const c10::Error& e) {
     std::cerr << "error loading the model\n";
     ASSERT_TRUE(false);
@@ -35,7 +35,7 @@ TEST(CppAPITest, ResNetModuleFallbacksCorrectly) {
 TEST(CppAPITest, MobileNetModuleFallbacksCorrectlyWithOneEngine) {
   torch::jit::script::Module mod;
   try {
-    mod = torch::jit::load("tests/modules/mobilenet_v2_traced.jit.pt");
+    mod = torch::jit::load("tests/modules/mobilenet_v2_scripted.jit.pt");
   } catch (const c10::Error& e) {
     std::cerr << "error loading the model\n";
     ASSERT_TRUE(false);

tests/cpp/test_modules_as_engines.cpp

@@ -29,7 +29,6 @@ TEST_P(CppAPITests, ModuleToEngineToModuleIsClose) {
   std::vector<at::Tensor> jit_results;
   jit_results.push_back(jit_results_ivalues.toTensor());
 
-  auto forward_graph = mod.get_method("forward");
   std::vector<c10::ArrayRef<int64_t>> input_ranges;
   for (auto in : inputs) {
     input_ranges.push_back(in.sizes());
@@ -43,7 +42,7 @@ TEST_P(CppAPITests, ModuleToEngineToModuleIsClose) {
   auto engine = trtorch::ConvertGraphToTRTEngine(mod, "forward", input_ranges);
   auto trt_mod = trtorch::EmbedEngineInNewModule(engine, compile_spec.device);
 
-  torch::jit::IValue trt_results_ivalues = trtorch::tests::util::RunModuleForward(mod, inputs_ivalues);
+  torch::jit::IValue trt_results_ivalues = trtorch::tests::util::RunModuleForward(trt_mod, inputs_ivalues);
   std::vector<at::Tensor> trt_results;
   trt_results.push_back(trt_results_ivalues.toTensor());
 
@@ -61,4 +60,4 @@ INSTANTIATE_TEST_SUITE_P(
         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-3})));
+        PathAndInSize({"tests/modules/vit_scripted.jit.pt", {{1, 3, 224, 224}}, 8e-3})));