def test_calls_step_callbacks(self):
    """Runs a training loop, asserting that callbacks are called."""
    call_at_steps = [1, 3, 4]
    begin_steps = []
    end_steps = []
    test_case = self

    class TestCallback(callbacks.TrainingStepCallback):

      def call_at(self, step):
        return step in call_at_steps

      def on_step_begin(self, step):
        begin_steps.append(step)

      def on_step_end(self, step):
        # Assert that on_step_begin() was called before.
        test_case.assertIn(step, begin_steps)
        end_steps.append(step)

    model = tl.Serial(tl.Dense(1))
    task = training.TrainTask(
        _very_simple_data(), tl.L2Loss(), optimizers.SGD(.01)
    )
    loop = training.Loop(model, [task], callbacks=[TestCallback])
    loop.run(n_steps=5)

    # Assert that the callback has been called at the appropriate steps.
    self.assertEqual(begin_steps, call_at_steps)
    self.assertEqual(end_steps, call_at_steps)
Runs a training loop, asserting that callbacks are called.
trax/supervised/training_test.py