def test_initializes_step_callbacks_with_loop_instance(self):
    """Runs a training loop, asserting that callbacks are initialized."""

    class ActualLoop:
      # Wrapper object to make the Loop reference mutable.
      loop = None

    class TestCallback(callbacks.TrainingStepCallback):

      def __init__(self, loop):
        super().__init__(loop)
        ActualLoop.loop = loop

      def call_at(self, step):
        return False

      def on_step_begin(self, step):
        del step

      def on_step_end(self, step):
        del step

    model = tl.Serial(tl.Dense(1))
    task = training.TrainTask(
        _very_simple_data(), tl.L2Loss(), optimizers.SGD(.01)
    )
    expected_loop = training.Loop(
        model, [task], callbacks=[TestCallback]
    )
    self.assertIs(ActualLoop.loop, expected_loop)
Runs a training loop, asserting that callbacks are initialized.
trax/supervised/training_test.py