fix nblocks issues (#432)

intel · Feb 11, 2025 · edc2b03 · edc2b03
1 parent 40ab982
commit edc2b03
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/auto_round/autoround.py b/auto_round/autoround.py
@@ -176,7 +176,6 @@ def __init__(
         self.enable_quanted_input = enable_quanted_input
         self.enable_minmax_tuning = enable_minmax_tuning
         self.nsamples = nsamples
-        self.nblocks = nblocks
         self.bits = bits
         self.enable_norm_bias_tuning = enable_norm_bias_tuning
         self.group_size = group_size
@@ -1324,15 +1323,15 @@ def quant_blocks(
 
         if pbar is None:
             pbar = tqdm(range(0, len(block_names), nblocks))
-        # for i in pbar:
-        for i in range(len(block_names)):
+
+        for i in range(0, len(block_names), nblocks):
             if nblocks == 1:
                 n = block_names[i]
                 pbar.set_description(f"Quantizing {n}")
                 m = get_module(model, n)
             else:
-                names = block_names[i: i + nblocks]
-                pbar.set_description(f"Quantizing [{i + 1}-{i + nblocks}]/{len(block_names)}")
+                names = block_names[i: min(i + nblocks, len(block_names))]
+                pbar.set_description(f"Quantizing [{i + 1}-{min(i + nblocks, len(block_names))}]/{len(block_names)}")
                 modules = [get_module(model, n) for n in names]
                 m = WrapperMultiblock(modules)