Merge pull request #501 from letian-jiang/enable_kgraph_in_python

Enable LogLevel.GRAPH in Python API
pytorch · Jul 28, 2021 · b99ebb2 · b99ebb2
2 parents bd72677 + d03b476
commit b99ebb2
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/py/trtorch/csrc/trtorch_py.cpp b/py/trtorch/csrc/trtorch_py.cpp
@@ -320,6 +320,7 @@ PYBIND11_MODULE(_C, m) {
       .value("WARNING", core::util::logging::LogLevel::kWARNING)
       .value("INFO", core::util::logging::LogLevel::kINFO)
       .value("DEBUG", core::util::logging::LogLevel::kDEBUG)
+      .value("GRAPH", core::util::logging::LogLevel::kGRAPH)
       .export_values();
 }
 

diff --git a/py/trtorch/logging.py b/py/trtorch/logging.py
@@ -13,6 +13,7 @@ class Level(Enum):
     Warning = LogLevel.WARNING
     Info = LogLevel.INFO
     Debug = LogLevel.DEBUG
+    Graph = LogLevel.GRAPH
 
     @staticmethod
     def _to_internal_level(external) -> LogLevel:
@@ -26,6 +27,8 @@ def _to_internal_level(external) -> LogLevel:
             return LogLevel.INFO
         if external == Level.Debug:
             return LogLevel.DEBUG
+        if external == Level.Graph:
+            return LogLevel.GRAPH
 
 
 def get_logging_prefix() -> str: