def test_meta_device_materialization():
    """Test that the `setup_module()` method materializes meta-device tensors in the module."""

    class NoResetParameters(nn.Module):
        def __init__(self):
            super().__init__()
            self.weight = nn.Parameter(torch.ones(4, 4))

    class CustomModel(nn.Module):
        def __init__(self):
            super().__init__()
            # nn.Sequential as a parameterless module
            self.layer1 = nn.Sequential(NoResetParameters(), NoResetParameters())
            self.layer2 = nn.Linear(4, 4)
            self.register_buffer("buffer", torch.rand(2))

        def reset_parameters(self):
            self.buffer.fill_(1.0)

    strategy = ModelParallelStrategy(parallelize_fn=(lambda x, _: x))
    strategy._device_mesh = Mock()
    strategy._parallel_devices = [torch.device("cpu")]

    with torch.device("meta"):
        model = CustomModel()
    assert model.layer1[0].weight.is_meta
    assert model.layer2.weight.is_meta
    assert model.buffer.is_meta

    with pytest.warns(UserWarning, match=r"`reset_parameters\(\)` method for re-initialization: NoResetParameters"):
        model = strategy.setup_module(model)
    assert all(not p.is_meta for p in model.parameters())
    assert all(not b.is_meta for b in model.buffers())
Test that the `setup_module()` method materializes meta-device tensors in the module.
tests/tests_fabric/strategies/test_model_parallel.py