code: def __call__(self, x: jnp.ndarray):
      # normalize
      x = x.astype(jnp.float32) / 255.0
  
      # base
      x = ConvBlock()(x, 32, (3, 3))
      x = ConvBlock()(x, 64, (3, 3), stride=2)
      x = ConvBlock()(x, 64, (3, 3), stride=2)
      x = ConvBlock()(x, 128, (3, 3), stride=2)
  
      # GlobalAveragePooling2D
      x = jnp.mean(x, axis=(1, 2))
  
      # 1x1 Conv
      x = eg.Linear(10)(x)
  
      return x
  
file_path: elegy/examples/elegy/mnist_torch_dataloader.py
imports: import os
  import typing
  from datetime import datetime
  import jax
  import jax.numpy
  import matplotlib.pyplot
  import numpy
  import optax
  import torch
  import typer
  from datasets.load import load_dataset
  from tensorboardX.writer import SummaryWriter
  from torch.utils.data import DataLoader, TensorDataset
  import elegy
label_desc: Use image node_id 1 for calling the CNN obj's underlying member method code with example usage: obj.__call__(x) and returns: x