fix(QModule): support modules without weights

huggingface · Oct 29, 2024 · 7aaf99e · 7aaf99e
1 parent d563d4d
commit 7aaf99e
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/optimum/quanto/nn/qmodule.py b/optimum/quanto/nn/qmodule.py
@@ -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)
@@ -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(