def test_module_init_context(distributed, precision, expected_dtype):
    """Test that the module under the init-context gets moved to the right device and dtype."""
    strategy = ModelParallelStrategy(parallelize_fn=_parallelize_feed_forward_fsdp2)
    fabric = Fabric(accelerator="cuda", devices=2, strategy=strategy, precision=precision)
    fabric.launch()

    def _run_setup_assertions(empty_init, expected_device):
        with fabric.init_module(empty_init=empty_init):
            model = FeedForward()

        # The model is on the CPU/meta-device until after `.setup()``
        assert all(weight.device == expected_device for weight in model.parameters())
        assert all(weight.dtype == expected_dtype for weight in model.parameters())
        model = fabric.setup(model)
        # Parameters get sharded in `.setup()` and moved to the target device
        assert all(weight.device == torch.device("cuda", fabric.local_rank) for weight in model.parameters())
        assert all(weight.dtype == expected_dtype for weight in model.parameters())

    _run_setup_assertions(empty_init=False, expected_device=torch.device("cpu"))
    _run_setup_assertions(empty_init=True, expected_device=torch.device("meta"))
Test that the module under the init-context gets moved to the right device and dtype.
tests/tests_fabric/strategies/test_model_parallel_integration.py