import torch


def load_checkpoint(checkpoint_path):
    """
    Loads the model's, optimizer's and the scheduler's state dicts for continuing training or later testing.

        Parameters:
        -----------
            checkpoint_path (string) : path to the .pt file checkpoint

        Returns:
        --------
            corresponding state dictionaries
    """
    
    state = torch.load(checkpoint_path)
    
    model_state = state['model_dict']
    optimizer_state = state['optim_dict']
    scheduler_state = state['scheduler_dict']
       
    return model_state, optimizer_state, scheduler_state