浏览代码

Fix removal of nonexistent keys

Max Ryabinin 1 年之前
父节点
当前提交
16fb547960
共有 1 个文件被更改,包括 2 次插入1 次删除
  1. 2 1
      src/petals/server/from_pretrained.py

+ 2 - 1
src/petals/server/from_pretrained.py

@@ -65,7 +65,8 @@ def load_pretrained_block(
 
     # dummy load, check that keys match
     report = block.load_state_dict(state_dict, strict=False)
-    report.missing_keys.pop("self_attn.qkv_proj.weight", None)  # will be filled later
+    if "self_attn.qkv_proj.weight" in report.missing_keys:
+        report.missing_keys.remove("self_attn.qkv_proj.weight")  # will be filled later
     assert not report.missing_keys, f"Some block weights are missing: {report.missing_keys}"
 
     for param_name, _ in block.named_parameters():