def test_train_dense_layer_with_momentum(self):
    """Trains with an optimizer that has slots / requires initialization."""
    model = tl.Serial(tl.Dense(1))
    task = training.TrainTask(
        _very_simple_data(), tl.L2Loss(), optimizers.Momentum(.01))
    eval_task = training.EvalTask(
        _very_simple_data(),  # deliberately re-using training data
        [tl.L2Loss()],
        metric_names=['Momentum.L2Loss'])
    training_session = training.Loop(model, [task], eval_tasks=[eval_task],
                                     eval_at=lambda step_n: step_n % 2 == 0)
    self.assertEqual(0, training_session.step)
    training_session.run(n_steps=20)
    self.assertEqual(20, training_session.step)
Trains with an optimizer that has slots / requires initialization.
trax/supervised/training_test.py