fix comment

huggingface · Oct 16, 2024 · 50a405a · 50a405a
1 parent 3b6ddfc
commit 50a405a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/optimum/gptq/quantizer.py b/optimum/gptq/quantizer.py
@@ -408,7 +408,7 @@ def quantize_model(self, model: nn.Module, tokenizer: Optional[Any] = None):
  blocks = recurse_getattr(model, self.block_name_to_quantize)
 
  if not has_device_map:
- # put modules from module_name_preceding_first_block on cuda
+ # put modules from module_name_preceding_first_block on cuda or cpu
  to_device = "cuda:0" if torch.cuda.is_available() else "cpu"
  for module_name in self.module_name_preceding_first_block:
  module = recurse_getattr(model, module_name)