def test_trains_on_two_tasks(self):
    """Trains a very simple network on two very simple tasks."""
    model = tl.Serial(tl.Dense(3), tl.Dense(1))
    task = training.TrainTask(
        _very_simple_data(),
        tl.L2Loss(),
        optimizers.SGD(.01)
    )
    eval_task = training.EvalTask(
        _very_simple_data(),  # deliberately re-using training data
        [tl.L2Loss()],
    )
    training_session = training.Loop(
        model,
        tasks=(task, task),
        eval_tasks=(eval_task, eval_task),
        which_task=lambda step_n: step_n % 2,
    )
    self.assertEqual(0, training_session.step)
    training_session.run(n_steps=15)
    self.assertEqual(15, training_session.step)
    training_session.run(n_steps=5)
    self.assertEqual(20, training_session.step)
Trains a very simple network on two very simple tasks.
trax/supervised/training_test.py