Fix error on 4bit checkpoint load with run_lm_eval on TF4.45.2 (#1439)

huggingface · Oct 18, 2024 · 190a29a · 190a29a
1 parent caf0a52
commit 190a29a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/text-generation/utils.py b/examples/text-generation/utils.py
@@ -691,7 +691,7 @@ def local_load_remaining_pretrained_weight(self,model):
             "model": model,
             "state_dict": state_dict,
             "start_prefix": "",
-            "expected_keys": list(state_dict.keys()),
+            "expected_keys": self.loaded_state_dict_keys,
             "device_map": {"": self.device},
             "offload_folder": offload_folder,
             "state_dict_folder": tempfile.mkdtemp() if offload_state_dict else None,