code: def ConvBlock(
      x,
      units: int,
      kernel: tp.Tuple[int, int],
      stride: int = 1,
  ):
      x = eg.Conv(
          units,
          kernel,
          strides=[stride, stride],
          padding="same",
      )(x)
      x = eg.BatchNorm()(x)
      x = eg.Dropout(0.2)(x)
      return jax.nn.relu(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 a global function with example usage: ConvBlock(x, units, kernel, stride) and returns: jax