code: def test_pred_step(self):
      N = 0
  
      class Model(ModelCore):
          a: jnp.ndarray = tx.node()
  
          def init_step(
              self,
              key: jnp.ndarray,
              inputs: tp.Any,
          ) -> "Model":
              self.a = jnp.array(0, dtype=jnp.int32)
              return self
  
          def pred_step(self, inputs):
              nonlocal N
              N += 1
  
              self.a += 1
  
              return 1, self
  
      model = Model()
  
      preds = model.predict_on_batch(inputs=np.array(1.0))
      assert N == 1
      assert preds == 1
      assert model.a == 1
  
      preds = model.predict_on_batch(inputs=np.array(1.0))
      assert N == 1
      assert preds == 1
      assert model.a == 2
  
      model.eager = True
  
      preds = model.predict_on_batch(inputs=(np.array(1.0)))
      assert N == 2
      assert preds == 1
      assert model.a == 3
  
file_path: elegy/tests/model/model_core_test.py
imports: import typing
  import unittest
  import jax
  import jax.numpy
  import numpy
  import treex
  import elegy
  from elegy.model.model_core import ModelCore
label_desc: Use image node_id 2 for calling the ModelCoreTest obj's underlying member method code with example usage: obj.test_pred_step() and returns: self, int,  self