def test_train_one_task_eval_two_tasks(self):
    """Trains a very simple network on one task and evaluates on two tasks."""
    model = tl.Serial(tl.Dense(3), tl.Dense(1))
    task = training.TrainTask(
        _very_simple_data(),
        tl.L2Loss(),
        optimizers.SGD(.01)
    )
    export_prefix_1 = 'eval_1'
    eval_task_1 = training.EvalTask(
        _very_simple_data(),  # deliberately re-using training data
        [tl.L2Loss()],
        export_prefix=export_prefix_1,
    )
    export_prefix_2 = 'eval_2'
    eval_task_2 = training.EvalTask(
        _very_simple_data(),  # deliberately re-using training data
        [tl.L2Loss()],
        export_prefix=export_prefix_2,
    )
    training_session = training.Loop(
        model,
        tasks=(task,),
        eval_tasks=(eval_task_1, eval_task_2),
    )
    self.assertEqual(0, training_session.step)
    training_session.run(n_steps=5)
    self.assertEqual(5, training_session.step)
    export_prefixes = [task.export_prefix
                       for task in training_session.eval_tasks]
    self.assertCountEqual([export_prefix_1, export_prefix_2],
                          export_prefixes)
Trains a very simple network on one task and evaluates on two tasks.
trax/supervised/training_test.py