diff --git a/quantization/image_classification/cpu/resnet50_data_reader.py b/quantization/image_classification/cpu/resnet50_data_reader.py index d07198f11..0909c3554 100644 --- a/quantization/image_classification/cpu/resnet50_data_reader.py +++ b/quantization/image_classification/cpu/resnet50_data_reader.py @@ -46,16 +46,16 @@ def __init__(self, calibration_image_folder: str, model_path: str): (_, _, height, width) = session.get_inputs()[0].shape # Convert image to input data - self.nhwc_data_list = _preprocess_images( + self.nchw_data_list = _preprocess_images( calibration_image_folder, height, width, size_limit=0 ) self.input_name = session.get_inputs()[0].name - self.datasize = len(self.nhwc_data_list) + self.datasize = len(self.nchw_data_list) def get_next(self): if self.enum_data is None: self.enum_data = iter( - [{self.input_name: nhwc_data} for nhwc_data in self.nhwc_data_list] + [{self.input_name: nchw_data} for nchw_data in self.nchw_data_list] ) return next(self.enum_data, None)