code: def evaluate(data_path, master_csv_path, split_path, save_dir):
      """Evaluation Script."""
      logging.info("Evaluating OGB molviv")
      logging.info("Dataset split: %s", split_path)
      # Initialize the dataset reader.
      reader = data_utils.DataReader(
          data_path=data_path,
          master_csv_path=master_csv_path,
          split_path=split_path,
          batch_size=1,
      )
      # Transform impure `net_fn` to pure functions with hk.transform.
      net = hk.without_apply_rng(hk.transform(net_fn))
      with pathlib.Path(save_dir, "molhiv.pkl").open("rb") as fp:
          params = pickle.load(fp)
      accumulated_loss = 0
      accumulated_accuracy = 0
      idx = 0
  
      # We jit the computation of our loss, since this is the main computation.
      # Using jax.jit means that we will use a single accelerator. If you want
      # to use more than 1 accelerator, use jax.pmap. More information can be
      # found in the jax documentation.
      compute_loss_fn = jax.jit(
          functools.partial(compute_loss, net=net)
      )
      for graph in reader:
          # Jax will re-jit your graphnet every time a new graph shape is encountered.
          # In the limit, this means a new compilation every training step, which
          # will result in *extremely* slow training. To prevent this, pad each
          # batch of graphs to the nearest power of two. Since jax maintains a cache
          # of compiled programs, the compilation cost is amortized.
          graph = pad_graph_to_nearest_power_of_two(graph)
  
          # Extract the labels and remove from the graph.
          label = graph.globals["label"]
          graph = graph._replace(globals={})
          loss, acc = compute_loss_fn(params, graph, label)
          accumulated_accuracy += acc
          accumulated_loss += loss
          idx += 1
          if idx % 100 == 0:
              logging.info("Evaluated %s graphs", idx)
      logging.info("Completed evaluation.")
      loss = accumulated_loss / idx
      accuracy = accumulated_accuracy / idx
      logging.info("Eval loss: %s, accuracy %s", loss, accuracy)
      return loss, accuracy
  
imports: import functools
  import logging
  import pathlib
  import pickle
  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 9 for calling a global function with example usage: evaluate(data_path, master_csv_path, split_path, save_dir) and returns: loss,  accuracy