import torch


def load_checkpoint(checkpoint_name, device='cpu'):
    """
    DESCRIPTION: The function loads a torchvision model along with all attributes.
    INPUT: <checkpoint_name>: Full path of the checkpoint file (incl. file extension)
           <device> (default='cpu'): The device on which the model should be loaded.
                                     Typically 'cuda' to make use of the GPU (fast) or 'cpu' (slow)
    OUTPUT: <model>: The model along with its state_dict as it was saved in the checkpoint file
            <optimizer>: The optimizer along with its state_dict as it was saved in the checkpoint file
            <epochs>: The number of epochs as it was saved in the checkpoint file
    """
    #01. Load the checkpoint
    if device == 'cuda':
        device += ':0'
    cp = torch.load(checkpoint_name, map_location=device)

    #02. Apply the checkpoint attributes
    model = cp['model']
    model.load_state_dict(cp['model_state_dict'])
    model.class_to_idx = cp['class_to_idx']
    optimizer = cp['optimizer']
    optimizer.load_state_dict(cp['optimizer_state_dict'])
    epochs = cp['epochs']

    #03. Return the updated results
    return model, optimizer, epochs