docstr: Train MNIST model a bit, to compare to other implementations.
docstr_tok: ["Train","MNIST","model","a","bit",",","to","compare","to","other","implementations","."]
code: 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
      )
  
file_path: trax/trax/supervised/mnist_test.py
imports: import io
  import itertools
  from unittest import mock
  from absl.testing import absltest
  from trax import layers
  from trax.data import inputs
  from trax.data import tf_inputs
  from trax.optimizers import adam
  from trax.supervised import training
label_desc: Use image node_id 1 for calling the MnistTest obj's underlying member method code with example usage: obj.test_train_mnist_single_task(mock_stdout) without return types