code: def compute_loss(params, graph, label, net):
      """Computes loss."""
      pred_graph = net.apply(params, graph)
      preds = jax.nn.log_softmax(pred_graph.globals)
      targets = jax.nn.one_hot(label, 2)
  
      # Since we have an extra 'dummy' graph in our batch due to padding, we want
      # to mask out any loss associated with the dummy graph.
      # Since we padded with `pad_with_graphs` we can recover the mask by using
      # get_graph_padding_mask.
      mask = jraph.get_graph_padding_mask(pred_graph)
  
      # Cross entropy loss.
      loss = -jnp.mean(preds * targets * mask[:, None])
  
      # Accuracy taking into account the mask.
      accuracy = jnp.sum(
          (jnp.argmax(pred_graph.globals, axis=1) == label) * mask
      ) / jnp.sum(mask)
      return loss, accuracy
  
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 6 for calling a global function with example usage: compute_loss(params, graph, label, net) and returns: loss,  accuracy