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)

    # Can use additional `assign` argument
    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
Test that state_dict access passes through to the original module.
tests/tests_fabric/test_wrappers.py