code: def test_fabric_module_state_dict_access():
      """Test that state_dict access passes through to the original module."""
  
      class OriginalModule(torch.nn.Module):
          def __init__(self):
              super().__init__()
              self.layer = torch.nn.Linear(2, 3)
  
      original_module = OriginalModule()
  
      class ModuleWrapper(torch.nn.Module):
          def __init__(self):
              super().__init__()
              self.wrapped = original_module
  
      wrapped_module = ModuleWrapper()
  
      fabric_module = _FabricModule(
          wrapped_module, Mock(), original_module=original_module
      )
      state_dict = fabric_module.state_dict()
      assert set(state_dict.keys()) == {"layer.weight", "layer.bias"}
  
      weight, bias = torch.rand(3, 2), torch.rand(3)
      fabric_module.load_state_dict(
          {"layer.weight": weight, "layer.bias": bias}
      )
      assert torch.equal(fabric_module.layer.weight, weight)
      assert torch.equal(fabric_module.layer.bias, bias)
  
      if _TORCH_GREATER_EQUAL_2_1:
          # Can use additional `assign` argument in PyTorch >= 2.1
          with torch.device("meta"):
              original_module = OriginalModule()
          fabric_module = _FabricModule(
              wrapped_module, Mock(), original_module=original_module
          )
          assert fabric_module.layer.weight.is_meta
          fabric_module.load_state_dict(
              {"layer.weight": weight, "layer.bias": bias}, assign=True
          )
          assert not fabric_module.layer.weight.is_meta
  
imports: from unittest import mock
  from unittest.mock import Mock, call
  import pytest
  import torch
  from lightning.fabric.fabric import Fabric
  from lightning.fabric.plugins import Precision
  from lightning.fabric.utilities.device_dtype_mixin import _DeviceDtypeModuleMixin
  from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_1
  from lightning.fabric.wrappers import _FabricDataLoader, _FabricModule, _FabricOptimizer, _unwrap_objects, is_wrapped
  from torch.utils.data import BatchSampler, DistributedSampler
  from torch.utils.data.dataloader import DataLoader
  from tests_fabric.helpers.runif import RunIf
label_desc: Use image node_id 5 for calling a global function with example usage: test_fabric_module_state_dict_access() without return types