check W_q in state_dict

mobiusml · Feb 20, 2025 · c60218e · c60218e
1 parent 0bd8d2f
commit c60218e
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/hqq/core/quantize.py b/hqq/core/quantize.py
@@ -680,11 +680,13 @@ def _load_from_state_dict(
                 layer_state_dict[key] = state_dict.pop(prefix + key)
             else:
                 if(key not in ['bias']):
-                    missing_keys.append(prefix + key)
-
-        layer_state_dict['W_q'] = nn.Parameter(layer_state_dict['W_q'], requires_grad=False)
+                    missing_keys.append(prefix + key)                    
 
-        self.load_state_dict(layer_state_dict, strict=strict)
+        if 'W_q' in layer_state_dict:
+            layer_state_dict['W_q'] = nn.Parameter(layer_state_dict['W_q'], requires_grad=False)
+            self.load_state_dict(layer_state_dict, strict=strict)
+        else:
+            missing_keys.append(prefix + "W_q")  
 
     def load_state_dict(self, state_dict, strict=True, assign=False):
         if "encoded_state_dict" in state_dict: