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

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>
pytorch · Oct 9, 2021 · 282e98a · 282e98a
1 parent a38f0c7
commit 282e98a
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/tests/cpp/test_modules_as_engines.cpp b/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})));