code: def load_dataset(dataset_name):
      """Load the dataset and return the graph, features, train/valid/test sets
      and the number of classes.
  
      Here, we use `BuiltInDataset` to load the dataset which returns graph,
      features, train/valid/test sets and the number of classes.
      """
      dataset = gb.BuiltinDataset(dataset_name).load()
      print(f"Loaded dataset: {dataset.tasks[0].metadata['name']}")
  
      graph = dataset.graph
      features = dataset.feature
      train_set = dataset.tasks[0].train_set
      valid_set = dataset.tasks[0].validation_set
      test_set = dataset.tasks[0].test_set
      num_classes = dataset.tasks[0].metadata["num_classes"]
  
      return (
          graph,
          features,
          train_set,
          valid_set,
          test_set,
          num_classes,
      )
  
imports: import argparse
  import itertools
  import sys
  import time
  import dgl
  import dgl.graphbolt
  import dgl.nn
  import psutil
  import torch
  import torch.nn
  import torch.nn.functional
  from dgl.nn import HeteroEmbedding
  from ogb.lsc import MAG240MEvaluator
  from ogb.nodeproppred import Evaluator
  from tqdm import tqdm
label_desc: Use image node_id 1 for calling a global function with example usage: load_dataset(dataset_name) and returns: graph,  features,  train_set,  valid_set,  test_set,  num_classes