class_docstr: Class with all code to run experiment.
class_docstr_tok: ["Class","with","all","code","to","run","experiment","."]
docstr: Computes loss function.
  
  Args:
    images: tensor with images NCHW
    labels: tensors with dense labels, shape (batch_size,)
  
  Returns:
    Tuple (total_loss, losses_dictionary).
docstr_tok: ["Computes","loss","function",".","Args",":","images",":","tensor","with","images","NCHW","labels",":","tensors","with","dense","labels",",","shape","(","batch_size",",",")","Returns",":","Tuple","(","total_loss",",","losses_dictionary",")","."]
code: def loss_fn(self, images, labels):
      """Computes loss function.
  
      Args:
        images: tensor with images NCHW
        labels: tensors with dense labels, shape (batch_size,)
  
      Returns:
        Tuple (total_loss, losses_dictionary).
      """
      logits = self.model(images, training=True)
      xent_loss = objax.functional.loss.cross_entropy_logits_sparse(
          logits, labels
      ).mean()
      wd_loss = (
          FLAGS.weight_decay
          * 0.5
          * sum(
              (v.value**2).sum()
              for k, v in self.model_vars.items()
              if k.endswith(".w")
          )
      )
      accuracy = (
          jn.count_nonzero(jn.equal(jn.argmax(logits, axis=1), labels))
          / (images.shape[0] + 1e-5)
          * 100
      )
      total_loss = xent_loss + wd_loss
      return total_loss, {
          "total_loss": total_loss,
          "xent_loss": xent_loss,
          "wd_loss": wd_loss,
          "accuracy": accuracy,
      }
  
file_path: objax/examples/image_classification/imagenet_resnet50_train.py
imports: import os
  import time
  import jax
  import jax.numpy
  import numpy
  from absl import app, flags, logging
  import objax
  from examples.image_classification import imagenet_resnet50_data
  from objax.zoo.resnet_v2 import ResNet50
label_desc: Use image node_id 4 for calling the Experiment obj's underlying member method code with example usage: obj.loss_fn(images, labels) and returns: total_loss,  dict