code: def test_test_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 test_step(
              self, inputs, labels
          ) -> elegy.TestStepOutput["Model"]:
              nonlocal N
              N += 1
              self.a += 1
              loss = 1.0
  
              return loss, dict(loss=loss), self
  
      model = Model()
  
      logs = model.test_on_batch(inputs=(np.array(1.0)), labels=(1.0,))
      assert N == 1
      assert logs["loss"] == 1.0
      assert model.a == 1
  
      logs = model.test_on_batch(inputs=(np.array(1.0)), labels=(1.0,))
      assert N == 1
      assert logs["loss"] == 1.0
      assert model.a == 2
  
      model.eager = True
  
      logs = model.test_on_batch(inputs=(np.array(1.0)), labels=(1.0,))
      assert N == 2
      assert logs["loss"] == 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 3 for calling the ModelCoreTest obj's underlying member method code with example usage: obj.test_test_step() and returns: self, loss,  dict,  self