def test_module_mode_restore_missing_module():
    """Test that restoring still works if the module drops a layer after it was captured."""

    class Model(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.child1 = torch.nn.Linear(2, 2)
            self.child2 = torch.nn.Linear(2, 2)

    model = Model()
    mode = _ModuleMode()
    mode.capture(model)
    model.child1.eval()
    del model.child2
    assert not hasattr(model, "child2")
    mode.restore(model)
    assert model.child1.training
Test that restoring still works if the module drops a layer after it was captured.
tests/tests_pytorch/utilities/test_model_helpers.py