import torch

'''Utility for Storing and Restoring models between epochs in Training'''

def save_checkpoint(file_name, model, optimizer, epochs_finished, other_info={}):

    checkpoint = {
        'model' :  model.state_dict(),
        'optimizer' : optimizer.state_dict(),
        'epochs_finished': epochs_finished,
        'other_info' : other_info # Additional Parameter Dict
    }

    torch.save(checkpoint, file_name)

    
def restore_checkpoint(file_name, model, optimizer):

    checkpoint = torch.load(file_name)
    model.load_state_dict(checkpoint['model'])
    optimizer.load_state_dict(checkpoint['optimizer'])
    current_epoch = checkpoint['epochs_finished'] + 1
    other_info = checkpoint['other_info']

    return model, optimizer, current_epoch, other_info

import torch

'''Utility for Storing and Restoring models between epochs in Training'''

def save_checkpoint(file_name, model, optimizer, epochs_finished, other_info={}):

    checkpoint = {
        'model' :  model.state_dict(),
        'optimizer' : optimizer.state_dict(),
        'epochs_finished': epochs_finished,
        'other_info' : other_info # Additional Parameter Dict
    }

    torch.save(checkpoint, file_name)

    
def restore_checkpoint(file_name, model, optimizer):

    checkpoint = torch.load(file_name)
    model.load_state_dict(checkpoint['model'])
    optimizer.load_state_dict(checkpoint['optimizer'])
    current_epoch = checkpoint['epochs_finished'] + 1
    other_info = checkpoint['other_info']

    return model, optimizer, current_epoch, other_info

Utility for Storing and Restoring models between epochs in Training # Additional Parameter Dict