Skip to content

Commit af20761

Browse files
committed
chore: linter fixes
Signed-off-by: Dheeraj Peri <[email protected]>
1 parent 749048c commit af20761

File tree

1 file changed

+6
-2
lines changed

1 file changed

+6
-2
lines changed

py/torch_tensorrt/ptq.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,12 +55,14 @@ def write_calibration_cache(self, cache):
5555
else:
5656
return b""
5757

58+
5859
# deepcopy (which involves pickling) is performed on the compile_spec internally during compilation.
5960
# We register this __reduce__ function for pickler to identity the calibrator object returned by DataLoaderCalibrator during deepcopy.
6061
# This should be the object's local name relative to the module https://docs.python.org/3/library/pickle.html#object.__reduce__
6162
def __reduce__(self):
6263
return self.__class__.__name__
6364

65+
6466
class DataLoaderCalibrator(object):
6567
"""
6668
Constructs a calibrator class in TensorRT and uses pytorch dataloader to load/preproces
@@ -119,7 +121,7 @@ def __new__(cls, *args, **kwargs):
119121
"get_batch": get_cache_mode_batch if use_cache else get_batch,
120122
"read_calibration_cache": read_calibration_cache,
121123
"write_calibration_cache": write_calibration_cache,
122-
"__reduce__": __reduce__ # used when you deepcopy the DataLoaderCalibrator object
124+
"__reduce__": __reduce__, # used when you deepcopy the DataLoaderCalibrator object
123125
}
124126

125127
# Using type metaclass to construct calibrator class based on algorithm type
@@ -129,7 +131,9 @@ def __new__(cls, *args, **kwargs):
129131
)()
130132
elif algo_type == CalibrationAlgo.ENTROPY_CALIBRATION_2:
131133
return type(
132-
"Int8EntropyCalibrator2", (_C.IInt8EntropyCalibrator2,), attribute_mapping
134+
"Int8EntropyCalibrator2",
135+
(_C.IInt8EntropyCalibrator2,),
136+
attribute_mapping,
133137
)()
134138
elif algo_type == CalibrationAlgo.LEGACY_CALIBRATION:
135139
return type(

0 commit comments

Comments
 (0)