(save_dir, save_name, model, optim, log)
| 51 | |
| 52 | |
| 53 | def save_checkpoint(save_dir, save_name, model, optim, log): |
| 54 | torch.save({ |
| 55 | 'model_state_dict': utils.get_model_state(model), |
| 56 | 'optim_state_dict': optim.state_dict(), |
| 57 | }, os.path.join(save_dir, '{}-model.pkl'.format(save_name))) |
| 58 | with open(os.path.join(save_dir, '{}-log.pkl'.format(save_name)), 'wb') as f: |
| 59 | pickle.dump(log, f) |
| 60 | |
| 61 | |
| 62 | def main(args, logger): |