only change model keys of key name starts with 'model'

unitaryai · Sep 3, 2021 · 5700981 · 5700981
1 parent 8d5b1cf
commit 5700981
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/convert_weights.py b/convert_weights.py
@@ -14,8 +14,9 @@ def main():
         "config": checkpoint["hyper_parameters"]["config"],
     }
     for k, v in checkpoint["state_dict"].items():
-        name = k[6:]  # remove `model.`
-        new_state_dict["state_dict"][name] = v
+        if k.startswith("model."):
+            k = k[6:]  # remove `model.`
+        new_state_dict["state_dict"][k] = v
 
     torch.save(new_state_dict, ARGS.save_to)