code: def _replace_label(graph):
      if cfg.dataset.augment_label:
          if not cfg.dataset.augment_label_dims:
              raise ValueError(
                  "dataset.label_dims should have the same "
                  "length as dataset.label"
              )
          label_key = _key(cfg.dataset.augment_label, as_label=True)
          label = graph[label_key]
          if cfg.dataset.task == "node":
              graph.node_label = label
          elif (
              cfg.dataset.task == "edge"
              or cfg.dataset.task == "link_pred"
          ):
              graph.edge_label = label
          elif cfg.dataset.task == "graph":
              graph.graph_label = label
          else:
              raise ValueError(
                  "Unknown task type: {}".format(cfg.dataset.task)
              )
  
imports: import logging
  import networkx
  import numpy
  import torch
  import torch.nn
  from deepsnap.graph import Graph
  import graphgym.register
  from graphgym.config import cfg
  from graphgym.contrib.transform.identity import compute_identity
label_desc: Use image node_id 3 for calling a global function with example usage: _replace_label(graph) without return types