Batch size

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)