code: def main():
      args = arg_parse()
  
      edge_train_mode = args.mode
      print("edge train mode: {}".format(edge_train_mode))
  
      WN_graph = nx.read_gpickle(args.data_path)
      print(
          "Each node has node ID (n_id). Example: ", WN_graph.nodes[0]
      )
      print(
          "Each edge has edge ID (id) and categorical label (e_label). Example: ",
          WN_graph[0][5871],
      )
  
      # Since both feature and label are relation types,
      # Only the disjoint mode would make sense
      node_feature = torch.ones(WN_graph.number_of_nodes(), 5)
      edge_index = torch.LongTensor(list(WN_graph.edges())).permute(
          1, 0
      )
      graph = Graph(
          node_feature=node_feature,
          edge_index=edge_index,
          directed=True,
      )
  
      dataset = GraphDataset(
          [graph],
          task="link_pred",
          edge_train_mode=edge_train_mode,
          edge_message_ratio=args.edge_message_ratio,
          edge_negative_sampling_ratio=args.neg_sampling_ratio,
      )
  
      # find num edge types
      max_label = 0
      labels = []
      for u, v, edge_key in WN_graph.edges:
          l = WN_graph[u][v][edge_key]["e_label"]
          if not l in labels:
              labels.append(l)
      # labels are consecutive (0-17)
      num_edge_types = len(labels)
  
      print("Pre-transform: ", dataset[0])
      dataset = dataset.apply_transform(
          WN_transform,
          update_tensor=False,
          G=WN_graph,
          num_edge_types=num_edge_types,
          deep_copy=False,
      )
      print("Post-transform: ", dataset[0])
      print(
          "Initial data: {} nodes; {} edges.".format(
              dataset[0].num_nodes, dataset[0].num_edges
          )
      )
      print(
          "Number of node features: {}".format(
              dataset.num_node_features
          )
      )
  
      # split dataset
      datasets = {}
      (
          datasets["train"],
          datasets["val"],
          datasets["test"],
      ) = dataset.split(transductive=True, split_ratio=[0.8, 0.1, 0.1])
  
      print("After split:")
      print(
          "Train message-passing graph: {} nodes; {} edges.".format(
              datasets["train"][0].num_nodes,
              datasets["train"][0].num_edges,
          )
      )
      print(
          "Val message-passing graph: {} nodes; {} edges.".format(
              datasets["val"][0].num_nodes, datasets["val"][0].num_edges
          )
      )
      print(
          "Test message-passing graph: {} nodes; {} edges.".format(
              datasets["test"][0].num_nodes,
              datasets["test"][0].num_edges,
          )
      )
  
      # node feature dimension
      input_dim = datasets["train"].num_node_features
      edge_feat_dim = datasets["train"].num_edge_features
      num_classes = datasets["train"].num_edge_labels
      print(
          "Node feature dim: {}; edge feature dim: {}; num classes: {}.".format(
              input_dim, edge_feat_dim, num_classes
          )
      )
  
      # relation type is both used for edge features and edge labels
      model = Net(input_dim, edge_feat_dim, num_classes, args).to(
          args.device
      )
      optimizer = torch.optim.Adam(
          model.parameters(), lr=0.001, weight_decay=5e-3
      )
      follow_batch = []  # e.g., follow_batch = ['edge_index']
  
      dataloaders = {
          split: DataLoader(
              ds,
              collate_fn=Batch.collate(follow_batch),
              batch_size=1,
              shuffle=(split == "train"),
          )
          for split, ds in datasets.items()
      }
      print("Graphs after split: ")
      for key, dataloader in dataloaders.items():
          for batch in dataloader:
              print(key, ": ", batch)
  
      train(model, dataloaders, optimizer, args)
  
imports: import copy
  import time
  import argparse
  import torch
  import torch.nn
  import torch.nn.functional
  import networkx
  import numpy
  import sklearn.metrics
  import torch_geometric.nn
  from torch.utils.data import DataLoader
  from deepsnap.dataset import GraphDataset
  from deepsnap.batch import Batch
  from deepsnap.graph import Graph
label_desc: Use image node_id 5 for calling a global function with example usage: main() without return types