code: def main(conf):
      conf.model.params = conf.model.params[conf.dataset.name]
      params = conf.model.params
      logger.info(OmegaConf.to_yaml(conf))
      try:
          edge_dropout = params.edge_dropout
      except:
          edge_dropout = 0.0
  
      grad_norm = (
          None
          if isinstance(params.grad_norm, str)
          else params.grad_norm
      )
  
      device = (
          f"cuda:{conf.device}" if torch.cuda.is_available() else "cpu"
      )
  
      t = time.perf_counter()
      logger.info("Loading data...")
      data, in_channels, out_channels = get_data(
          conf.root, conf.dataset.name
      )
      logger.info(f"Done! [{time.perf_counter() - t:.2f}s]")
      perm, ptr = metis(
          data.adj_t, num_parts=params.num_parts, log=True
      )
      data = permute(data, perm, log=True)
  
      if conf.model.loop:
          t = time.perf_counter()
          logger.info("Adding self-loops...")
          data.adj_t = data.adj_t.set_diag()
          logger.info(f"Done! [{time.perf_counter() - t:.2f}s]")
      if conf.model.norm:
          t = time.perf_counter()
          logger.info("Normalizing data...")
          data.adj_t = gcn_norm(data.adj_t, add_self_loops=False)
          logger.info(f"Done! [{time.perf_counter() - t:.2f}s]")
  
      if data.y.dim() == 1:
          criterion = torch.nn.CrossEntropyLoss()
      else:
          criterion = torch.nn.BCEWithLogitsLoss()
  
      train_loader = SubgraphLoader(
          data,
          ptr,
          batch_size=params.batch_size,
          shuffle=True,
          num_workers=params.num_workers,
          persistent_workers=params.num_workers > 0,
      )
  
      eval_data = data
      eval_loader = EvalSubgraphLoader(
          eval_data, ptr, batch_size=params["batch_size"]
      )
  
      t = time.perf_counter()
      logger.info("Calculating buffer size...")
      # We reserve a much larger buffer size than what is actually needed for
      # training in order to perform efficient history accesses during inference.
  
      # set buffer_size to None, which makes the model._async=False
      # Here, we implemented without offset and count in the forward method to disable the model._async
      buffer_size = max(
          [n_id.numel() for _, _, n_id, _, _ in eval_loader]
      )
      logger.info(
          f"Done! [{time.perf_counter() - t:.2f}s] -> {buffer_size}"
      )
  
      kwargs = {}
      if conf.model.name[:3] == "PNA":
          kwargs["deg"] = data.adj_t.storage.rowcount()
  
      GNN = getattr(models, conf.model.name)
      # we want to mask the model._async=False
      model = GNN(
          num_nodes=data.num_nodes,
          in_channels=in_channels,
          out_channels=out_channels,
          pool_size=conf.model.params.pool_size,
          buffer_size=buffer_size,
          gamma=conf.model.params.gamma,
          **params.architecture,
          **kwargs,
      ).to(device)
  
      optimizer = torch.optim.Adam(
          [
              dict(
                  params=model.reg_modules.parameters(),
                  weight_decay=params.reg_weight_decay,
              ),
              dict(
                  params=model.nonreg_modules.parameters(),
                  weight_decay=params.nonreg_weight_decay,
              ),
          ],
          lr=params.lr,
      )
  
      t = time.perf_counter()
      logger.info("Fill history...")
      mini_test(model, eval_loader)
      logger.info(f"Done! [{time.perf_counter() - t:.2f}s]")
  
      best_val_acc = test_acc = 0
      for epoch in range(1, params.epochs + 1):
          loss = mini_train(
              model,
              train_loader,
              criterion,
              optimizer,
              params.max_steps,
              grad_norm,
              edge_dropout,
          )
          out = mini_test(model, eval_loader)
          train_acc = compute_micro_f1(out, data.y, data.train_mask)
  
          val_acc = compute_micro_f1(out, data.y, data.val_mask)
          tmp_test_acc = compute_micro_f1(out, data.y, data.test_mask)
  
          if val_acc > best_val_acc:
              best_val_acc = val_acc
              test_acc = tmp_test_acc
          if epoch % conf.log_every == 0:
              logger.info(
                  f"Epoch: {epoch:04d}, Loss: {loss:.4f}, "
                  f"Train: {train_acc:.4f}, Val: {val_acc:.4f}, "
                  f"Test: {tmp_test_acc:.4f}, Final: {test_acc:.4f}"
              )
  
      logger.info("=========================")
      logger.info(f"Val: {best_val_acc:.4f}, Test: {test_acc:.4f}")
  
imports: import time
  import hydra
  from omegaconf import OmegaConf
  import torch
  from torch_geometric.nn.conv.gcn_conv import gcn_norm
  from dig.lsgraph.dataset import get_data, SubgraphLoader, EvalSubgraphLoader
  from dig.lsgraph.method.GraphFMOB.metis import metis, permute
  import dig.lsgraph.method.GraphFMOB.models
  from dig.lsgraph.method.GraphFMOB.utils import compute_micro_f1, dropout
  from torch import Tensor
  from typing import Optional
  from torch_geometric.typing import Adj
  from torch_geometric.utils.num_nodes import maybe_num_nodes
  from torch_scatter import scatter_add
  from torch_sparse import SparseTensor
  import logging
label_desc: Use image node_id 5 for calling a global function with example usage: main(conf) without return types