Skip to content

Commit

Permalink
fix(QModule): support modules without weights
Browse files Browse the repository at this point in the history
  • Loading branch information
dacorvo committed Oct 29, 2024
1 parent 500e77d commit 97a6266
Showing 1 changed file with 3 additions and 2 deletions.
5 changes: 3 additions & 2 deletions optimum/quanto/nn/qmodule.py
Original file line number Diff line number Diff line change
Expand Up @@ -217,7 +217,8 @@ def from_module(
if qmodule is None:
return None
# Move the quantized module to the target device, but with empty weights
qmodule = qmodule.to_empty(device=module.weight.device)
device = torch.device("cpu") if module.weight is None else module.weight.device
qmodule = qmodule.to_empty(device=device)
# Set scales that were initialized to empty values
qmodule.input_scale = torch.ones_like(qmodule.input_scale)
qmodule.output_scale = torch.ones_like(qmodule.output_scale)
Expand All @@ -226,7 +227,7 @@ def from_module(
if module.bias is not None:
qmodule.bias = module.bias

return qmodule.to(module.weight.device)
return qmodule.to(device)

@classmethod
def qcreate(
Expand Down

0 comments on commit 97a6266

Please sign in to comment.