Skip to content

Commit dae72d4

Browse files
authored
fix cli usage
1 parent a6bbed1 commit dae72d4

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

auto_round/compressors/base.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -625,7 +625,8 @@ def _set_auto_device_map_in_block(self, block: torch.nn.Module, input_ids: list[
625625
if self.low_gpu_mem_usage:
626626
input_output_memory = 0
627627

628-
mem_per_param_scale = 13 * self.batch_size if self.mem_per_param_scale is None else self.mem_per_param_scale
628+
mem_per_param_scale = 13 if self.mem_per_param_scale is None else self.mem_per_param_scale
629+
mem_per_param_scale *= self.batch_size
629630
if self.iters == 0:
630631
mem_per_param_scale = 1 # for rtn
631632

0 commit comments

Comments
 (0)