def test_train_eval_mode_restored(tmp_path):
    """Test that the trainer restores the `training` mode of all submodules to what it was before entering the loop."""

    class MixedTrainModeModule(BoringModel):
        def __init__(self):
            super().__init__()
            # A frozen submodule should keep its mode, regardless of whether we're training or not
            self.frozen = torch.nn.Linear(2, 2)
            self.frozen.eval()
            self.frozen.requires_grad_(False)

        def training_step(self, *args, **kwargs):
            assert self.layer.weight.requires_grad
            assert self.layer.training
            assert not self.frozen.training
            assert not self.frozen.weight.requires_grad
            return super().training_step(*args, **kwargs)

        def validation_step(self, *args, **kwargs):
            assert self.layer.weight.requires_grad
            assert not self.layer.training
            assert not self.frozen.training
            assert not self.frozen.weight.requires_grad
            return super().validation_step(*args, **kwargs)

        def test_step(self, *args, **kwargs):
            assert self.layer.weight.requires_grad
            assert not self.layer.training
            assert not self.frozen.training
            assert not self.frozen.weight.requires_grad
            return super().test_step(*args, **kwargs)

        def predict_step(self, *args, **kwargs):
            assert self.layer.weight.requires_grad
            assert not self.layer.training
            assert not self.frozen.training
            assert not self.frozen.weight.requires_grad
            return super().predict_step(*args, **kwargs)

    model = MixedTrainModeModule()
    trainer = Trainer(
        devices=1,
        default_root_dir=tmp_path,
        max_epochs=1,
        val_check_interval=1,
        limit_train_batches=3,
        limit_val_batches=2,
        limit_test_batches=2,
        limit_predict_batches=2,
        enable_progress_bar=False,
        enable_model_summary=False,
        enable_checkpointing=False,
    )
    trainer.fit(model)
    trainer.validate(model)
    trainer.test(model)
    trainer.predict(model)
Test that the trainer restores the `training` mode of all submodules to what it was before entering the loop.
tests/tests_pytorch/models/test_hooks.py