Enable LogLevel.GRAPH in Python API

Signed-off-by: Letian Jiang <[email protected]>
pytorch · Jun 17, 2021 · d03b476 · d03b476
1 parent eb39f9c
commit d03b476
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
@@ -309,6 +309,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: