code: def loss(x, label):  # sum(label * log(softmax(logit)))
      logit = model(x)
      return objax.functional.loss.cross_entropy_logits(
          logit, label
      ).mean()
  
imports: import collections
  import random
  import re
  import jax
  import jax.numpy
  import tensorflow_datasets
  import objax
  from objax.functional import one_hot
  from objax.zoo.rnn import RNN
label_desc: Use image node_id 7 for calling a global function with example usage: loss(x, label) and returns: objax