code: def create_augment_fun(graph_fun, key):
      def augment_fun(graph, as_label, **kwargs):
          graph[_key(key, as_label=as_label)] = graph_fun(
              graph, **kwargs
          )
  
      return augment_fun
  
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 2 for calling a global function with example usage: create_augment_fun(graph_fun, key) and returns: augment_fun