class_docstr: CNN model.
class_docstr_tok: ["CNN","model","."]
code: def compute_loss(
      self, predictions: Predictions, input_batch: NestedMap
  ) -> tuple[WeightedScalars, NestedMap]:
      predictions = cast(NestedMap, predictions)
      loss = _cross_entropy_loss(input_batch["labels"], predictions)
      accuracy = _compute_accuracy(input_batch["labels"], predictions)
      return NestedMap(
          loss=(loss, 1.0), accuracy=(accuracy, 1.0)
      ), NestedMap(predictions=predictions)
  
imports: from typing import cast
  import jax
  import jax.numpy
  from paxml import base_experiment
  from paxml import base_task
  from paxml import experiment_registry
  from paxml import learners
  from paxml import tasks_lib
  from praxis import base_input
  from praxis import base_layer
  from praxis import base_model
  from praxis import optimizers
  from praxis import pax_fiddle
  from praxis import py_utils
  from praxis import pytypes
  from praxis import schedules
  from praxis.layers import activations
  from praxis.layers import convolutions
  from praxis.layers import linears
  from praxis.layers import poolings
  from pax_examples.dali_pax_input import MnistDaliInput
label_desc: Use image node_id 3 for calling the CNNModel obj's underlying member method code with example usage: obj.compute_loss(predictions, input_batch) and returns: NestedMap,  NestedMap