chore(//py/torch_tensorrt): ptq mypy compliance

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

Author: narendasan

py/torch_tensorrt/ptq.py

@@ -1,4 +1,4 @@
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Self, Optional
 import torch
 import os
 
@@ -16,15 +16,15 @@ class CalibrationAlgo(Enum):
     MINMAX_CALIBRATION = _C.CalibrationAlgo.MINMAX_CALIBRATION
 
 
-def get_cache_mode_batch(self):
+def get_cache_mode_batch(self: object) -> None:
     return None
 
 
-def get_batch_size(self):
+def get_batch_size(self: object) -> int:
     return 1
 
 
-def get_batch(self, names):
+def get_batch(self: object, _: Any) -> Optional[List[int]]:
     if self.current_batch_idx + self.batch_size > len(self.data_loader.dataset):
         return None
 
@@ -39,27 +39,30 @@ def get_batch(self, names):
     return inputs_gpu
 
 
-def read_calibration_cache(self):
+def read_calibration_cache(self: object) -> bytes:
     if self.cache_file and self.use_cache:
         if os.path.exists(self.cache_file):
             with open(self.cache_file, "rb") as f:
-                return f.read()
+                b: bytes = f.read()
+                return b
+        else:
+            raise FileNotFoundError(self.cache_file)
     else:
         return b""
 
 
-def write_calibration_cache(self, cache):
+def write_calibration_cache(self: object, cache: bytes) -> None:
     if self.cache_file:
         with open(self.cache_file, "wb") as f:
             f.write(cache)
     else:
-        return b""
+        return
 
 
 # deepcopy (which involves pickling) is performed on the compile_spec internally during compilation.
 # We register this __reduce__ function for pickler to identity the calibrator object returned by DataLoaderCalibrator during deepcopy.
 # This should be the object's local name relative to the module https://docs.python.org/3/library/pickle.html#object.__reduce__
-def __reduce__(self):
+def __reduce__(self: object) -> str:
     return self.__class__.__name__
 
 
@@ -75,10 +78,10 @@ class DataLoaderCalibrator(object):
         device: device on which calibration data is copied to.
     """
 
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs: Any):
         pass
 
-    def __new__(cls, *args, **kwargs):
+    def __new__(cls: Self, *args: Any, **kwargs: Any) -> Self:
         dataloader = args[0]
         algo_type = kwargs.get("algo_type", CalibrationAlgo.ENTROPY_CALIBRATION_2)
         cache_file = kwargs.get("cache_file", None)
@@ -158,10 +161,10 @@ class CacheCalibrator(object):
         algo_type: choice of calibration algorithm.
     """
 
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs: Any):
         pass
 
-    def __new__(cls, *args, **kwargs):
+    def __new__(cls: Self, *args: Any, **kwargs: Any) -> Self:
         cache_file = args[0]
         algo_type = kwargs.get("algo_type", CalibrationAlgo.ENTROPY_CALIBRATION_2)