def load_data(dataset):
    r"""Load knowledge graph dataset for RGCN link prediction tasks

    It supports three datasets: wn18, FB15k and FB15k-237

    Parameters
    ----------
    dataset: str
        The name of the dataset to load.

    Return
    ------
    The dataset object.
    """
    if dataset == "wn18":
        return WN18Dataset()
    elif dataset == "FB15k":
        return FB15kDataset()
    elif dataset == "FB15k-237":
        return FB15k237Dataset()
Load knowledge graph dataset for RGCN link prediction tasks

    It supports three datasets: wn18, FB15k and FB15k-237

    Parameters
    ----------
    dataset: str
        The name of the dataset to load.

    Return
    ------
    The dataset object.
    
python/dgl/data/knowledge_graph.py