code: def loss(x, label):
      logit = model(x)
      return objax.functional.loss.cross_entropy_logits(
          logit, label
      ).mean()
  
file_path: objax/examples/image_classification/mnist_dnn.py
imports: import os
  import numpy
  import tensorflow_datasets
  from tqdm import trange
  import objax
  from objax.functional import leaky_relu, one_hot
  from objax.jaxboard import SummaryWriter, Summary
  from objax.util import EasyDict
  from objax.zoo.dnnet import DNNet
label_desc: Use image node_id 1 for calling a global function with example usage: loss(x, label) and returns: objax