code: def main():
      args = arg_parse()
  
      edge_train_mode = args.mode
      print("edge train mode: {}".format(edge_train_mode))
  
      G = nx.read_gpickle(args.data_path)
      print(G.number_of_edges())
      print("Each node has node ID (n_id). Example: ", G.nodes[0])
      print(
          "Each edge has edge ID (id) and categorical label (e_label). Example: ",
          G[0][5871],
      )
  
      # find num edge types
      max_label = 0
      labels = []
      for u, v, edge_key in G.edges:
          l = G[u][v][edge_key]["e_label"]
          if not l in labels:
              labels.append(l)
      # labels are consecutive (0-17)
      num_edge_types = len(labels)
  
      H = WN_transform(G, num_edge_types)
      # The nodes in the graph have the features: node_feature and node_type (just one node type "n1" here)
      for node in H.nodes(data=True):
          print(node)
          break
      # The edges in the graph have the features: edge_feature and edge_type ("0" - "17" here)
      for edge in H.edges(data=True):
          print(edge)
          break
  
      hetero = HeteroGraph(H)
  
      if edge_train_mode == "disjoint":
          dataset = GraphDataset(
              [hetero],
              task="link_pred",
              edge_train_mode=edge_train_mode,
              edge_message_ratio=args.edge_message_ratio,
          )
      else:
          dataset = GraphDataset(
              [hetero],
              task="link_pred",
              edge_train_mode=edge_train_mode,
          )
  
      dataset_train, dataset_val, dataset_test = dataset.split(
          transductive=True, split_ratio=[0.8, 0.1, 0.1]
      )
      train_loader = DataLoader(
          dataset_train, collate_fn=Batch.collate(), batch_size=1
      )
      val_loader = DataLoader(
          dataset_val, collate_fn=Batch.collate(), batch_size=1
      )
      test_loader = DataLoader(
          dataset_test, collate_fn=Batch.collate(), batch_size=1
      )
      dataloaders = {
          "train": train_loader,
          "val": val_loader,
          "test": test_loader,
      }
  
      hidden_size = args.hidden_dim
      conv1, conv2 = generate_2convs_link_pred_layers(
          hetero, HeteroSAGEConv, hidden_size
      )
      model = HeteroGNN(conv1, conv2, hetero, hidden_size).to(
          args.device
      )
      optimizer = torch.optim.Adam(
          model.parameters(), lr=args.lr, weight_decay=args.weight_decay
      )
  
      t_accu, v_accu, e_accu = train(
          model, dataloaders, optimizer, args
      )
  
imports: import argparse
  import copy
  import time
  import numpy
  import networkx
  import sklearn.metrics
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.utils.data import DataLoader
  from torch_geometric.datasets import Planetoid
  from torch_geometric.datasets import TUDataset
  import torch_geometric.transforms
  import torch_geometric.nn
  from deepsnap.hetero_graph import HeteroGraph
  from deepsnap.dataset import GraphDataset
  from deepsnap.batch import Batch
  from deepsnap.hetero_gnn import HeteroSAGEConv, HeteroConv, forward_op
label_desc: Use image node_id 6 for calling a global function with example usage: main() without return types