def train_step(self, states, initializing):
return elegy.TrainStep(
logs=dict(loss=2.0),
states=elegy.States(
optimizer_states=0) if initializing else elegy.States(
optimizer_states=states.optimizer_states + 1),
)
def train_step(self, x, optimizer_states, initializing):
if initializing:
states = elegy.States(optimizer_states=0)
else:
states = elegy.States(optimizer_states=optimizer_states +
1)
return elegy.TrainStep(
logs=dict(loss=jnp.sum(x)),
states=states,
)
def train_step(self):
_, logs, states = self.test_step()
return elegy.TrainStep(logs, states.update(optimizer_states=4))