chore(//py/torch_tensorrt): __init__.py mypy conforming

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

Author: narendasan

py/torch_tensorrt/__init__.py

@@ -12,21 +12,23 @@
     __tensorrt_version__,
 )
 
+from typing import Dict, List
+
 if sys.version_info < (3,):
     raise Exception(
         "Python 2 has reached end-of-life and is not supported by Torch-TensorRT"
     )
 
 
-def _parse_semver(version):
+def _parse_semver(version: str) -> Dict[str, str]:
     split = version.split(".")
     if len(split) < 3:
         split.append("")
 
     return {"major": split[0], "minor": split[1], "patch": split[2]}
 
 
-def _find_lib(name, paths):
+def _find_lib(name: str, paths: List[str]) -> str:
     for path in paths:
         libpath = os.path.join(path, name)
         if os.path.isfile(libpath):
@@ -99,7 +101,7 @@ def _find_lib(name, paths):
     from torch_tensorrt.dynamo import backend
 
 
-def _register_with_torch():
+def _register_with_torch() -> None:
     trtorch_dir = os.path.dirname(__file__)
     torch.ops.load_library(trtorch_dir + "/lib/libtorchtrt.so")
 

py/torch_tensorrt/_util.py

@@ -4,7 +4,7 @@
 import torch
 
 
-def dump_build_info():
+def dump_build_info() -> None:
     """Prints build information about the torch_tensorrt distribution to stdout"""
     print(get_build_info())
 
@@ -15,20 +15,20 @@ def get_build_info() -> str:
     Returns:
         str: String containing the build information for torch_tensorrt distribution
     """
-    build_info = _C.get_build_info()
-    build_info = (
+    core_build_info = _C.get_build_info()
+    build_info = str(
         "Torch-TensorRT Version: "
         + str(__version__)
         + "\n"
         + "Using PyTorch Version: "
         + str(torch.__version__)
         + "\n"
-        + build_info
+        + core_build_info
     )
     return build_info
 
 
-def set_device(gpu_id):
+def set_device(gpu_id: int) -> None:
     _C.set_device(gpu_id)