Skip to content

Commit

Permalink
fix device check
Browse files Browse the repository at this point in the history
Signed-off-by: jiqing-feng <[email protected]>
  • Loading branch information
jiqing-feng committed Dec 24, 2024
1 parent d21256c commit 35fd5de
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions optimum/gptq/quantizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -519,7 +519,7 @@ def quantize_model(self, model: nn.Module, tokenizer: Optional[Any] = None):
blocks = recurse_getattr(model, self.block_name_to_quantize)

cur_layer_device = get_device(blocks[0])
if not is_gptqmodel_available():
if not is_gptqmodel_available() and cur_layer_device.type == "cpu":
cur_layer_device = 0

if not has_device_map:
Expand Down Expand Up @@ -591,7 +591,7 @@ def store_input_hook(_, input, *args):
block = block.to(0)
layers = get_layers(block)
block_device = get_device(block)
if not is_gptqmodel_available():
if not is_gptqmodel_available() and block_device.type == "cpu":
block_device = 0
if isinstance(self.modules_in_block_to_quantize, list) and len(self.modules_in_block_to_quantize) > 0:
if self.true_sequential:
Expand Down

0 comments on commit 35fd5de

Please sign in to comment.