(path, epoch, model, optimizer=None)
| 127 | |
| 128 | |
| 129 | def save_model(path, epoch, model, optimizer=None): |
| 130 | if isinstance(model, torch.nn.DataParallel): |
| 131 | state_dict = model.module.state_dict() |
| 132 | else: |
| 133 | state_dict = model.state_dict() |
| 134 | data = {'epoch': epoch, |
| 135 | 'state_dict': state_dict} |
| 136 | if not (optimizer is None): |
| 137 | data['optimizer'] = optimizer.state_dict() |
| 138 | torch.save(data, path) |