fix(//core): Do not compile hidden methods

(methods prefixed with _) Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
pytorch · Jun 5, 2020 · 6bd1a3f · 6bd1a3f
1 parent dcb1474
commit 6bd1a3f
Showing 1 changed file with 10 additions and 7 deletions.
diff --git a/core/compiler.cpp b/core/compiler.cpp
@@ -150,13 +150,16 @@ torch::jit::script::Module CompileGraph(const torch::jit::script::Module& mod,
     torch::jit::script::Module new_mod(mod._ivalue()->name() + "_trt");
     std::vector<std::shared_ptr<torch::jit::Graph>> graphs;
     for (const torch::jit::script::Method& method : mod.get_methods()) {
-        auto engine = ConvertGraphToTRTEngine(mod, method.name(), cfg);
-        auto new_g = std::make_shared<torch::jit::Graph>();
-        AddEngineToGraph(new_mod, new_g, engine);
-        auto new_method = new_mod._ivalue()->compilation_unit()->create_function(method.name(), new_g);
-        auto schema = GenerateGraphSchema(new_mod, new_method->name(), new_g);
-        new_mod.type()->addMethod(new_method);
-        new_method->setSchema(schema);
+        // Don't convert hidden methods
+        if (method.name().rfind("_", 0)) {
+            auto engine = ConvertGraphToTRTEngine(mod, method.name(), cfg);
+            auto new_g = std::make_shared<torch::jit::Graph>();
+            AddEngineToGraph(new_mod, new_g, engine);
+            auto new_method = new_mod._ivalue()->compilation_unit()->create_function(method.name(), new_g);
+            auto schema = GenerateGraphSchema(new_mod, new_method->name(), new_g);
+            new_mod.type()->addMethod(new_method);
+            new_method->setSchema(schema);
+        }
     }
 
     return new_mod;