def test_dp_module_state_dict():
    """Test that the module state dict gets retrieved without the prefixed wrapper keys from DP."""

    class DataParallelMock(MagicMock):
        def __instancecheck__(self, instance):
            # to make the strategy's `isinstance(model, DataParallel)` pass with a mock as class
            return True

    strategy = DataParallelStrategy(parallel_devices=[torch.device("cpu"), torch.device("cpu")])

    # Without DP applied (no setup call)
    original_module = torch.nn.Linear(2, 3)
    assert strategy.get_module_state_dict(original_module).keys() == original_module.state_dict().keys()

    # With DP applied (setup called)
    with mock.patch("lightning.fabric.strategies.dp.DataParallel", DataParallelMock):
        wrapped_module = strategy.setup_module(original_module)
        assert strategy.get_module_state_dict(wrapped_module).keys() == original_module.state_dict().keys()
Test that the module state dict gets retrieved without the prefixed wrapper keys from DP.
tests/tests_fabric/strategies/test_dp.py