code: def __call__(
      self, x, training=False
  ):  # x = (batch, colors, height, width)
      y = self.pre_conv(x)
      y = self.block1(y)
      y = self.block2(y)
      logits = self.post_conv(y).mean(
          (2, 3)
      )  # logits = (batch, nclass)
      if training:
          return logits
      return objax.functional.softmax(logits)
  
file_path: objax/examples/image_classification/mnist_cnn.py
imports: import os
  import numpy
  import tensorflow_datasets
  from tqdm import trange
  import objax
  from objax.util import EasyDict
label_desc: Use image node_id 2 for calling the SimpleNet obj's underlying member method code with example usage: obj.__call__(x, training) and returns: objax, logits