def test_setup_module_move_to_device(fabric_module_mock, move_to_device, distributed):
    """Test that `move_to_device` does nothing, ModelParallel decides which device parameters get moved to which device
    (sharding)."""
    from torch.distributed._tensor import DTensor

    strategy = ModelParallelStrategy(parallelize_fn=_parallelize_feed_forward_fsdp2)
    fabric = Fabric(accelerator="cuda", devices=2, strategy=strategy)
    fabric.launch()

    model = FeedForward()
    fabric_model = fabric.setup_module(model, move_to_device=move_to_device)
    fabric_module_mock.assert_not_called()

    # the linear layer got sharded and each part is on the expected device
    assert fabric_model.w1.weight.device == torch.device("cuda", fabric.local_rank)
    assert isinstance(fabric_model.w1.weight, DTensor)

    # 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)
Test that `move_to_device` does nothing, ModelParallel decides which device parameters get moved to which device
    (sharding).
tests/tests_fabric/strategies/test_model_parallel_integration.py