def test_module_mode_restore_new_module(caplog):
    """Test that restoring ignores newly added submodules after the module was captured."""

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

    model = Model()
    mode = _ModuleMode()
    mode.capture(model)
    model.child.eval()
    model.new_child = torch.nn.Linear(2, 2)
    with caplog.at_level(logging.DEBUG, logger="lightning.pytorch.utilities.model_helpers"):
        mode.restore(model)
    assert "Restoring training mode on module 'new_child' not possible" in caplog.text
Test that restoring ignores newly added submodules after the module was captured.
tests/tests_pytorch/utilities/test_model_helpers.py