class_docstr: CNN model.
class_docstr_tok: ["CNN","model","."]
code: def compute_predictions(self, input_batch: NestedMap) -> Predictions:
      return self.network(input_batch["inputs"])
  
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 2 for calling the CNNModel obj's underlying member method code with example usage: obj.compute_predictions(input_batch) and returns: self