diff --git a/hpvm/test/epoch_dnn/torch_dnn/quantizer.py b/hpvm/test/epoch_dnn/torch_dnn/quantizer.py index 05b2a2535b418ad5b2b98aea4d922927f49dafca..a473c0a057a5015f55064af8a3392a6b12935d47 100644 --- a/hpvm/test/epoch_dnn/torch_dnn/quantizer.py +++ b/hpvm/test/epoch_dnn/torch_dnn/quantizer.py @@ -40,10 +40,11 @@ def quantize( strat: str = "NONE", working_dir: PathLike = ".", output_name: str = "calib.txt", + eval_batchsize: int = 128 ): # possible quant strats ['NONE', 'AVG', 'N_STD', 'GAUSS', 'LAPLACE'] print("Quantizing...") - dataloader = DataLoader(dataset, batch_size=1) + dataloader = DataLoader(dataset, batch_size=eval_batchsize) # Collect Pre Quantization Stats distiller.utils.assign_layer_fq_names(model)