(model_g, model_c1, model_c2, save_path)
| 50 | |
| 51 | |
| 52 | def save_model(model_g, model_c1, model_c2, save_path): |
| 53 | save_dic = { |
| 54 | 'g_state_dict': model_g.state_dict(), |
| 55 | 'c1_state_dict': model_c1.state_dict(), |
| 56 | 'c2_state_dict': model_c2.state_dict(), |
| 57 | } |
| 58 | torch.save(save_dic, save_path) |
| 59 | |
| 60 | def load_model(model_g, model_c, load_path): |
| 61 | checkpoint = torch.load(load_path) |