code: 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()
          )
  
file_path: pytorch-lightning/tests/tests_fabric/strategies/test_dp.py
imports: from unittest import mock
  from unittest.mock import MagicMock, Mock
  import pytest
  import torch
  from lightning.fabric import Fabric
  from lightning.fabric.strategies import DataParallelStrategy
  from tests_fabric.helpers.runif import RunIf
  from tests_fabric.strategies.test_single_device import _run_test_clip_gradients
label_desc: Use image node_id 5 for calling a global function with example usage: test_dp_module_state_dict() and returns: True