def test_restores_memory_efficient_from_standard(self):
    """Training restores step from directory where it saved it."""
    self.skipTest('Broken by https://github.com/google/jax/pull/11234')
    model = tl.Serial(tl.Dense(4), tl.Dense(1))
    task_std = training.TrainTask(
        _very_simple_data(), tl.L2Loss(), optimizers.Adam(.0001))
    tmp_dir = self.create_tempdir().full_path
    loop = training.Loop(model, [task_std],
                         checkpoint_at=lambda step_n: step_n % 2 == 0,
                         output_dir=tmp_dir)
    loop.run(4)
    task_memeff = training.TrainTask(
        _very_simple_data(), tl.L2Loss(), optimizers.Adam)
    loop2 = training.Loop(model, [task_memeff], output_dir=tmp_dir,
                          use_memory_efficient_trainer=True)
    loop2.run(2)
    self.assertEqual(6, loop2.step)
Training restores step from directory where it saved it.
trax/supervised/training_test.py