code: def train(
      data_path,
      master_csv_path,
      split_path,
      batch_size,
      num_training_steps,
      save_dir,
  ):
      """OGB Training Script."""
      # Initialize the dataset reader.
      reader = data_utils.DataReader(
          data_path=data_path,
          master_csv_path=master_csv_path,
          split_path=split_path,
          batch_size=batch_size,
          dynamically_batch=True,
      )
      # Repeat the dataset forever for training.
      reader.repeat()
  
      # Transform impure `net_fn` to pure functions with hk.transform.
      net = hk.without_apply_rng(hk.transform(net_fn))
      # Get a candidate graph and label to initialize the network.
      graph = reader.get_graph_by_idx(0)
  
      # Initialize the network.
      logging.info("Initializing network.")
      params = net.init(jax.random.PRNGKey(42), graph)
      # Because we are training with multiple devices, params needs to have a
      # device axis.
      params = jax.device_put_replicated(params, list(jax.devices()))
      # Initialize the optimizer.
      opt_init, opt_update = optax.adam(1e-4)
      opt_state = jax.pmap(opt_init)(params)
  
      compute_loss_fn = functools.partial(compute_loss, net=net)
  
      # We pmap the computation of our loss, since this is the main computation.
      # Using jax.pmap means that we will use all available accelerators.
      # More information can be found in the jax documentation.
      @functools.partial(jax.pmap, axis_name="device")
      def update_fn(params, graph, label, opt_state):
          (loss, acc), grad = jax.value_and_grad(
              compute_loss_fn, has_aux=True
          )(params, graph, label)
          # Average gradients across devices
          grad = jax.lax.pmean(grad, axis_name="device")
          updates, opt_state = opt_update(grad, opt_state, params)
          params = optax.apply_updates(params, updates)
          return loss, acc, opt_state, params
  
      for idx in range(num_training_steps):
          graph_batch = next(device_batch(reader))
          label = graph_batch.globals["label"]
          loss, acc, opt_state, params = update_fn(
              params, graph_batch, label, opt_state
          )
          if idx % 100 == 0:
              logging.info(
                  "step: %s, loss: %s, acc: %s", idx, loss, acc
              )
      if save_dir is not None:
          with pathlib.Path(save_dir, "molhiv.pkl").open("wb") as fp:
              logging.info("Saving model to %s", save_dir)
              pickle.dump(params, fp)
      logging.info("Training finished")
  
imports: import functools
  import logging
  import pathlib
  import pickle
  from typing import Iterator
  from absl import app
  from absl import flags
  import haiku
  import jax
  import jax.numpy
  import jraph
  from jraph.ogb_examples import data_utils
  import optax
label_desc: Use image node_id 7 for calling a global function with example usage: train(data_path, master_csv_path, split_path, batch_size, num_training_steps, save_dir) and returns: loss,  acc,  opt_state,  params