code: def loss(x, y):
      logits = model(x, training=True)
      loss_xe = objax.functional.loss.cross_entropy_logits_sparse(
          logits, y
      ).mean()
      loss_l2 = 0.5 * sum(
          (v.value**2).sum()
          for k, v in model.vars().items()
          if k.endswith(".w")
      )
      return loss_xe + weight_decay * loss_l2, {
          "loss/xe": loss_xe,
          "loss/l2": loss_l2,
      }
  
file_path: objax/examples/image_classification/mnist_cnn.py
imports: import os
  import numpy
  import tensorflow_datasets
  from tqdm import trange
  import objax
  from objax.util import EasyDict
label_desc: Use image node_id 3 for calling a global function with example usage: loss(x, y) and returns: unknown,  dict