def test_loop_checkpoint_high_metric(self):
    """Runs a training loop that saves checkpoints for high metric values."""
    model = tl.Serial(tl.Dense(1))
    task = training.TrainTask(_very_simple_data(),
                              tl.L2Loss(),
                              optimizers.SGD(.01))
    eval_metric = tl.L2Loss()
    eval_task = training.EvalTask(_very_simple_data(),
                                  [eval_metric],
                                  metric_names=['l2_loss'])
    tmp_dir = self.create_tempdir().full_path
    loop = training.Loop(model,
                         [task],
                         eval_tasks=[eval_task],
                         output_dir=tmp_dir,
                         eval_at=lambda step_n: step_n % 2 == 0,
                         checkpoint_at=lambda step_n: step_n % 2 == 0,
                         checkpoint_high_metric='l2_loss')
    call_counter = collections.Counter()
    loop.save_checkpoint = lambda name: call_counter.update([name])
    loop.run(n_steps=10)

    # Eval metric steadily descends, so high checkpoint triggered only once.
    # Low checkpoint not defined, so never triggered.
    self.assertEqual(call_counter['model'], 5)
    self.assertEqual(call_counter['lowest_l2_loss'], 0)
    self.assertEqual(call_counter['highest_l2_loss'], 1)
Runs a training loop that saves checkpoints for high metric values.
trax/supervised/training_test.py