chore(//py): Adding documentation to the new context managers

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

py/torch_tensorrt/logging.py

@@ -108,6 +108,13 @@ def log(level: Level, msg: str):
     Graph = LogLevel.GRAPH
 
 class internal_errors:
+    """Context-manager to limit displayed log messages to just internal errors
+
+    Example::
+
+    with torch_tensorrt.logging.internal_errors():
+        outputs = model_torchtrt(inputs)
+    """
     def __enter__(self):
         self.external_lvl = get_reportable_log_level()
         set_reportable_log_level(Level.InternalError)
@@ -116,6 +123,13 @@ def __exit__(self, exc_type, exc_value, exc_tb):
         set_reportable_log_level(self.external_lvl)
 
 class errors:
+    """Context-manager to limit displayed log messages to just errors and above
+
+    Example::
+
+    with torch_tensorrt.logging.errors():
+        outputs = model_torchtrt(inputs)
+    """
     def __enter__(self):
         self.external_lvl = get_reportable_log_level()
         set_reportable_log_level(Level.Error)
@@ -124,6 +138,14 @@ def __exit__(self, exc_type, exc_value, exc_tb):
         set_reportable_log_level(self.external_lvl)
 
 class warnings:
+    """Context-manager to limit displayed log messages to just warnings and above
+
+    Example::
+
+    with torch_tensorrt.logging.warnings():
+        model_trt = torch_tensorrt.compile(model, **spec)
+    """
+
     def __enter__(self):
         self.external_lvl = get_reportable_log_level()
         set_reportable_log_level(Level.Warning)
@@ -132,6 +154,13 @@ def __exit__(self, exc_type, exc_value, exc_tb):
         set_reportable_log_level(self.external_lvl)
 
 class info:
+    """Context-manager to display all info and greater severity messages
+
+    Example::
+
+    with torch_tensorrt.logging.info():
+        model_trt = torch_tensorrt.compile(model, **spec)
+    """
     def __enter__(self):
         self.external_lvl = get_reportable_log_level()
         set_reportable_log_level(Level.Info)
@@ -140,6 +169,13 @@ def __exit__(self, exc_type, exc_value, exc_tb):
         set_reportable_log_level(self.external_lvl)
 
 class debug:
+    """Context-manager to display full debug information through the logger
+
+    Example::
+
+    with torch_tensorrt.logging.debug():
+        model_trt = torch_tensorrt.compile(model, **spec)
+    """
     def __enter__(self):
         self.external_lvl = get_reportable_log_level()
         set_reportable_log_level(Level.Debug)
@@ -148,6 +184,14 @@ def __exit__(self, exc_type, exc_value, exc_tb):
         set_reportable_log_level(self.external_lvl)
 
 class graphs:
+    """Context-manager to display the results of intermediate lowering passes
+    as well as full debug information through the logger
+
+    Example::
+
+    with torch_tensorrt.logging.graphs():
+        model_trt = torch_tensorrt.compile(model, **spec)
+    """
     def __enter__(self):
         self.external_lvl = get_reportable_log_level()
         set_reportable_log_level(Level.Graph)