Strip unwanted prefix from state keys when loading model

jmuletalbiach · Jan 5, 2023 · 2b9e168 · 2b9e168
1 parent ab04701
commit 2b9e168
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/sample.py b/sample.py
@@ -29,7 +29,12 @@
 checkpoint = torch.load(ckpt_path, map_location=device)
 gptconf = GPTConfig(**checkpoint['model_args'])
 model = GPT(gptconf)
-model.load_state_dict(checkpoint['model'])
+state_dict = checkpoint['model']
+unwanted_prefix = '_orig_mod.'
+for k,v in list(state_dict.items()):
+    if k.startswith(unwanted_prefix):
+        state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
+model.load_state_dict(state_dict)
 model.eval()
 model.to(device)
 if compile: