refactor: Fix the module for MLF test

narendasan · narendasan · commit 10cae301b1f0 · 2022-05-19T19:49:11.000-07:00
Signed-off-by: Naren Dasan &lt;naren@narendasan.com&gt;
Signed-off-by: Naren Dasan &lt;narens@nvidia.com&gt;
diff --git a/tests/core/lowering/test_module_fallback_passes.cpp b/tests/core/lowering/test_module_fallback_passes.cpp
@@ -18,7 +18,8 @@ TEST(Lowering, NotateModuleForFallbackWorksCorrectly) {
   }
 
   std::unordered_set<std::string> mods_to_mark;
-  mods_to_mark.insert("ModuleFallbackSub");
+  mods_to_mark.insert("custom_models.ModuleFallbackSub");
+
 
   torch_tensorrt::core::lowering::passes::NotateModuleForFallback(mod, "", "forward", mods_to_mark);
 
@@ -58,7 +59,7 @@ TEST(Lowering, MarkNodesForFallbackWorksCorrectly) {
   }
 
   std::unordered_set<std::string> mods_to_mark;
-  mods_to_mark.insert("ModuleFallbackSub");
+  mods_to_mark.insert("custom_models.ModuleFallbackSub");
 
   torch_tensorrt::core::lowering::passes::NotateModuleForFallback(mod, "", "forward", mods_to_mark);
   auto mod_ = torch::jit::freeze_module(mod);

Original file line number	Diff line number	Diff line change
`@@ -18,7 +18,8 @@ TEST(Lowering, NotateModuleForFallbackWorksCorrectly) {`
`18`	`18`	`}`
`19`	`19`
`20`	`20`	`std::unordered_set<std::string> mods_to_mark;`
`21`		`- mods_to_mark.insert("ModuleFallbackSub");`
	`21`	`+ mods_to_mark.insert("custom_models.ModuleFallbackSub");`
	`22`	`+`
`22`	`23`
`23`	`24`	`torch_tensorrt::core::lowering::passes::NotateModuleForFallback(mod, "", "forward", mods_to_mark);`
`24`	`25`
`@@ -58,7 +59,7 @@ TEST(Lowering, MarkNodesForFallbackWorksCorrectly) {`
`58`	`59`	`}`
`59`	`60`
`60`	`61`	`std::unordered_set<std::string> mods_to_mark;`
`61`		`- mods_to_mark.insert("ModuleFallbackSub");`
	`62`	`+ mods_to_mark.insert("custom_models.ModuleFallbackSub");`
`62`	`63`
`63`	`64`	`torch_tensorrt::core::lowering::passes::NotateModuleForFallback(mod, "", "forward", mods_to_mark);`
`64`	`65`	`auto mod_ = torch::jit::freeze_module(mod);`