code: 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()
  
imports: from __future__ import absolute_import
  import os, sys
  import pickle
  import networkx
  import numpy
  import scipy.sparse
  from ..None import backend
  from ..convert import graph
  from ..utils import retry_method_with_fix
  from .dgl_dataset import DGLBuiltinDataset
  from .utils import _get_dgl_url, deprecate_function, deprecate_property, download, extract_archive, generate_mask_tensor, get_download_dir, load_graphs, load_info, makedirs, save_graphs, save_info
label_desc: Use image node_id 5 for calling a global function with example usage: load_data(dataset) and returns: WN18Dataset, FB15kDataset, FB15k237Dataset