class_docstr: Model provides an Estimator-like API similar to Keras.
code: def __call__(self, *args, **kwargs) -> tp.Any:
      assert self.module is not None
  
      return self.module(*args, **kwargs)
  
imports: import typing
  import flax
  import jax
  import jax.numpy
  import numpy
  import treex
  from optax import GradientTransformation
  from elegy import types, utils
  from elegy.model.model_base import ModelBase
  from elegy.model.model_core import GradStepOutput, PredStepOutput, TestStepOutput, TrainStepOutput
label_desc: Use image node_id 4 for calling the Model obj's underlying member method code with example usage: obj.__call__() and returns: self