def test_step(self, metrics_states, initializing):
return elegy.TestStep(
loss=0.1,
logs=dict(loss=1.0),
states=elegy.States(metrics_states=0) if initializing else
elegy.States(metrics_states=metrics_states + 1),
)
def test_step(self, x, metrics_states, initializing):
if initializing:
states = elegy.States(metrics_states=0)
else:
states = elegy.States(metrics_states=metrics_states + 1)
return elegy.TestStep(
loss=0.1,
logs=dict(loss=jnp.sum(x)),
states=states,
)
def test_step(self):
_, states, _, _, _ = self.pred_step()
return elegy.TestStep(0, {}, states.update(metrics_states=3))