diff --git a/hpvm/test/epoch_dnn/torch_dnn/quantizer.py b/hpvm/test/epoch_dnn/torch_dnn/quantizer.py
index a473c0a057a5015f55064af8a3392a6b12935d47..ed94754d6a1bf09459c870dcd2b1cd58661c1a51 100644
--- a/hpvm/test/epoch_dnn/torch_dnn/quantizer.py
+++ b/hpvm/test/epoch_dnn/torch_dnn/quantizer.py
@@ -54,7 +54,7 @@ def quantize(
     if not os.path.isfile(stats_file):
         # generates `stats_file`
         collect_quant_stats(
-            model, lambda model: get_loss(model, dataloader), save_dir=working_dir
+            model, get_loss, dataloader, save_dir=working_dir
         )
 
     # Generate Quantized Scales
@@ -142,7 +142,7 @@ def get_loss(model: nn.Module, dataloader: DataLoader):
 
 
 @torch.no_grad()
-def evaluate(model: nn.Module, dataloader: DataLoader):
+def evaluate(model: nn.Module, dataloader: DataLoader = None):
     model.eval()
     correct = 0
     total = 0