def test_restore_fails_different_model(self):
    """Training restores from a checkpoint created with smaller model."""
    model1 = tl.Serial(tl.Dense(1))
    task = training.TrainTask(
        _very_simple_data(), tl.L2Loss(), optimizers.SGD(.01))
    tmp_dir = self.create_tempdir().full_path
    loop = training.Loop(model1, [task],
                         checkpoint_at=lambda step_n: step_n % 2 == 0,
                         output_dir=tmp_dir)
    loop.run(2)
    model2 = tl.Serial(tl.Dense(2))
    with self.assertRaises(IndexError):
      training.Loop(model2, [task], output_dir=tmp_dir)
Training restores from a checkpoint created with smaller model.
trax/supervised/training_test.py