| import argparse | |
| import torch | |
| def main(checkpoint): | |
| state_dict = torch.load(checkpoint, map_location="cpu") | |
| if "optimizer" in state_dict: | |
| del state_dict["optimizer"] | |
| if "lr_scheduler" in state_dict: | |
| del state_dict["lr_scheduler"] | |
| torch.save(state_dict, checkpoint) | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("checkpoint", type=str) | |
| args = parser.parse_args() | |
| main(args.checkpoint) | |