code: def test_module_mode():
      class ChildChildModule(torch.nn.Module):
          def __init__(self):
              super().__init__()
              self.layer = torch.nn.Linear(2, 2)
              self.dropout = torch.nn.Dropout()
  
      class ChildModule(torch.nn.Module):
          def __init__(self):
              super().__init__()
              self.child = ChildChildModule()
              self.dropout = torch.nn.Dropout()
  
      class RootModule(torch.nn.Module):
          def __init__(self):
              super().__init__()
              self.child1 = ChildModule()
              self.child2 = ChildModule()
              self.norm = torch.nn.BatchNorm1d(2)
  
      # Model with all submodules in the same mode
      model = RootModule()
      model.train()
      mode = _ModuleMode()
      mode.capture(model)
      model.eval()
      assert all(not m.training for m in model.modules())
      mode.restore(model)
      assert model.training
      assert all(m.training for m in model.modules())
      model.eval()
  
      mode = _ModuleMode()
      mode.capture(model)
      model.eval()
      assert all(not m.training for m in model.modules())
      mode.restore(model)
      assert all(not m.training for m in model.modules())
      model.train()
  
      # Model with submodules in different modes
      model.norm.eval()
      model.child1.eval()
      model.child2.train()
      model.child2.child.eval()
      model.child2.child.layer.train()
  
      mode = _ModuleMode()
      mode.capture(model)
      model.eval()
      assert all(not m.training for m in model.modules())
      mode.restore(model)
      assert model.training
      assert not model.norm.training
      assert all(not m.training for m in model.child1.modules())
      assert model.child2.training
      assert model.child2.dropout.training
      assert not model.child2.child.training
      assert model.child2.child.layer.training
      assert not model.child2.child.dropout.training
  
file_path: pytorch-lightning/tests/tests_pytorch/utilities/test_model_helpers.py
imports: import logging
  import pytest
  import torch.nn
  from lightning.pytorch import LightningDataModule
  from lightning.pytorch.demos.boring_classes import BoringDataModule, BoringModel
  from lightning.pytorch.utilities.model_helpers import _ModuleMode, _restricted_classmethod, is_overridden
  from lightning_utilities import module_available
label_desc: Use image node_id 4 for calling a global function with example usage: test_module_mode() without return types