fix: correct serialization of non-affine QLayerNorm

huggingface · Feb 14, 2025 · 1801dbc · 1801dbc
1 parent 2202f84
commit 1801dbc
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/optimum/quanto/nn/qmodule.py b/optimum/quanto/nn/qmodule.py
@@ -147,7 +147,9 @@ def disable_output_quantization(self):
     def _save_to_state_dict(self, destination, prefix, keep_vars):
         if self.weight_qtype is None or not self.frozen:
             # Save standard weight Tensor
-            destination[prefix + "weight"] = self.weight if keep_vars else self.weight.detach()
+            destination[prefix + "weight"] = (
+                self.weight if (self.weight is None or keep_vars) else self.weight.detach()
+            )
         else:
             # Save QTensor using dedicated method
             self.weight.save_to_state_dict(destination, prefix + "weight.", keep_vars)