Skip to content

Commit 7fd24b4

Browse files
committed
add progress bar for weight MSE calibration
Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
1 parent d686ac9 commit 7fd24b4

1 file changed

Lines changed: 3 additions & 1 deletion

File tree

modelopt/torch/quantization/model_calib.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -361,7 +361,9 @@ def mse_calibrate(
361361

362362
# Step 3: Calibrate weight quantizers ONE AT A TIME with immediate amax computation
363363
# This prevents massive memory accumulation seen in large models
364-
for idx, (parent_module, weight_name, weight_quantizer) in enumerate(weight_quantizers):
364+
for idx, (parent_module, weight_name, weight_quantizer) in enumerate(
365+
tqdm(weight_quantizers, desc="MSE weight calibration")
366+
):
365367
# Enable calibration mode for the weight quantizer
366368
weight_quantizer.disable_quant()
367369
weight_quantizer.enable_calib()

0 commit comments

Comments
 (0)