Fix removal of nonexistent keys

pull/498/head
Max Ryabinin 9 months ago
parent f100915641
commit 16fb547960

@ -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():

Loading…
Cancel
Save