code: def test_setup_module_move_to_device(
      fabric_module_mock, move_to_device
  ):
      """Test that `move_to_device` does nothing, FSDP decides which device parameters get moved to which device
      (sharding)."""
      strategy = FSDPStrategy(auto_wrap_policy=always_wrap_policy)
      fabric = Fabric(accelerator="cuda", devices=2, strategy=strategy)
      fabric.launch()
  
      model = torch.nn.Linear(
          10, 10, bias=False
      )  # total params: 10 * 10 = 100
      fabric_model = fabric.setup_module(
          model, move_to_device=move_to_device
      )
      fabric_module_mock.assert_not_called()
  
      assert len(list(fabric_model.parameters())) == 1
      # the linear layer got sharded and each part is on the expected device
      assert next(fabric_model.parameters()).device == torch.device(
          "cuda", fabric.local_rank
      )
      assert next(fabric_model.parameters()).numel() == 50
      if _TORCH_GREATER_EQUAL_2_0:
          # In PyTorch >= 2.0 we set `use_orig_params=True` and don't see flattened parameters
          assert isinstance(next(fabric_model.parameters()), Parameter)
      else:
          assert isinstance(
              next(fabric_model.parameters()), FlatParameter
          )
  
      # The _DeviceDtypeModuleMixin currently can't represent the device in a meaningful way for models with pieces on
      # different devices
      assert fabric_model.device == torch.device(
          "cuda", fabric.local_rank
      )
      assert fabric.device == torch.device("cuda", fabric.local_rank)
  
imports: import os
  from copy import deepcopy
  from pathlib import Path
  from unittest import mock
  import pytest
  import torch
  from lightning.fabric import Fabric
  from lightning.fabric.plugins import FSDPPrecision
  from lightning.fabric.strategies import FSDPStrategy
  from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_0, _TORCH_GREATER_EQUAL_2_1
  from lightning.fabric.wrappers import _FabricOptimizer
  from torch.distributed.fsdp import FlatParameter, FullyShardedDataParallel, OptimStateKeyType
  from torch.distributed.fsdp.wrap import always_wrap_policy, wrap
  from torch.nn import Parameter
  from tests_fabric.helpers.models import BoringFabric
  from tests_fabric.helpers.runif import RunIf
  from tests_fabric.test_fabric import BoringModel
label_desc: Use image node_id 4 for calling a global function with example usage: test_setup_module_move_to_device(fabric_module_mock, move_to_device) without return types