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

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

    class CustomModel(LightningModule):
        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)

        def configure_model(self) -> None:
            pass

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

    strategy = ModelParallelStrategy()
    strategy._accelerator = Mock()
    strategy._device_mesh = Mock()
    strategy._parallel_devices = [torch.device("cpu")]
    strategy._lightning_module = model
    strategy.model = model

    with pytest.warns(UserWarning, match=r"`reset_parameters\(\)` method for re-initialization: NoResetParameters"):
        strategy.setup(Mock())
    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()` method materializes meta-device tensors in the LightningModule.
tests/tests_pytorch/strategies/test_model_parallel.py