import torch


def save_checkpoint(train_data, epochs, model, path):
    """ Function to save the model checkpoint
        Checkpoint keys:
            - epochs (epoch number)
            - state_dict (model.state_dict())
            -class_to_idx (model.class_to_idx())
        
        Args:
            - epochs (int) 
            - model  
            - path (str ex. "/checkpoint.pth"
        
        Returns:
            - Model checkpoint file (extension .pth)
    """
    # move model to cpu to avoid loading issues
    device = torch.device("cpu") 
    model.to(device) 
    
    model.class_to_idx = train_data.class_to_idx
    
    checkpoint = {'epochs': epochs,
                  'state_dict': model.state_dict(),
                  'class_to_idx': model.class_to_idx}

    return torch.save(checkpoint, path)