@@ -55,12 +55,14 @@ def write_calibration_cache(self, cache):
55
55
else :
56
56
return b""
57
57
58
+
58
59
# deepcopy (which involves pickling) is performed on the compile_spec internally during compilation.
59
60
# We register this __reduce__ function for pickler to identity the calibrator object returned by DataLoaderCalibrator during deepcopy.
60
61
# This should be the object's local name relative to the module https://docs.python.org/3/library/pickle.html#object.__reduce__
61
62
def __reduce__ (self ):
62
63
return self .__class__ .__name__
63
64
65
+
64
66
class DataLoaderCalibrator (object ):
65
67
"""
66
68
Constructs a calibrator class in TensorRT and uses pytorch dataloader to load/preproces
@@ -119,7 +121,7 @@ def __new__(cls, *args, **kwargs):
119
121
"get_batch" : get_cache_mode_batch if use_cache else get_batch ,
120
122
"read_calibration_cache" : read_calibration_cache ,
121
123
"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
123
125
}
124
126
125
127
# Using type metaclass to construct calibrator class based on algorithm type
@@ -129,7 +131,9 @@ def __new__(cls, *args, **kwargs):
129
131
)()
130
132
elif algo_type == CalibrationAlgo .ENTROPY_CALIBRATION_2 :
131
133
return type (
132
- "Int8EntropyCalibrator2" , (_C .IInt8EntropyCalibrator2 ,), attribute_mapping
134
+ "Int8EntropyCalibrator2" ,
135
+ (_C .IInt8EntropyCalibrator2 ,),
136
+ attribute_mapping ,
133
137
)()
134
138
elif algo_type == CalibrationAlgo .LEGACY_CALIBRATION :
135
139
return type (
0 commit comments