osmarks/nanogpt-experiments
1
0
Fork 0
mirror of https://github.com/osmarks/nanogpt-experiments.git synced 2024-12-18 14:10:28 +00:00
Code Issues Releases Wiki Activity

Strip unwanted prefix from state keys when loading model

Browse Source
This commit is contained in:
Nat Friedman 2023-01-04 16:34:00 -08:00
parent ab04701f9f
commit 2b9e168736
1 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
sample.py
View File

@ -29,7 +29,12 @@ ckpt_path = os.path.join(out_dir, 'ckpt.pt')
checkpoint = torch.load(ckpt_path, map_location=device) checkpoint = torch.load(ckpt_path, map_location=device)
gptconf = GPTConfig(**checkpoint['model_args']) gptconf = GPTConfig(**checkpoint['model_args'])
model = GPT(gptconf) 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.eval()
model.to(device) model.to(device)
if compile: if compile: