def test_train_mnist_single_task(self, mock_stdout):
    """Train MNIST model a bit, to compare to other implementations."""
    mnist_model = _build_model(two_heads=False)
    (task, eval_task) = _mnist_tasks()
    training_session = training.Loop(
        mnist_model,
        tasks=[task],
        eval_tasks=[eval_task],
        eval_at=lambda step_n: step_n % 20 == 0,
    )

    training_session.run(n_steps=100)
    self.assertEqual(training_session.step, 100)

    # Assert that we reach at least 80% eval accuracy.
    self.assertGreater(_read_metric('WeightedCategoryAccuracy', mock_stdout),
                       0.8)
Train MNIST model a bit, to compare to other implementations.
trax/supervised/mnist_test.py