| import torch | |
| # Load the checkpoint | |
| checkpoint = torch.load('checkpoint_cleaned.pt', map_location=torch.device('cpu')) | |
| print(checkpoint.keys()) | |
| # remove keys of fform lang_encoder.gpt_neox.layers.x.decoder_layer | |
| for key in list(checkpoint.keys()): | |
| if 'decoder_layer' in key: | |
| del checkpoint[key] | |
| # save the checkpoint | |
| torch.save(checkpoint, 'checkpoint_cleaned.pt') |