import torch


def load_io_dataset(dataset_path: str, device=None):
    """
    Load the saved dataset, map the data location to `device`.

    Args:
        dataset_path: Path to the checkpoint (.pt format)
        device: Device for the data.

    Returns: The loaded IO dataset

    """
    data = torch.load(dataset_path, map_location=device)
    return data