def test_has_meta_device_parameters_or_buffers():
    """Test that the `_has_meta_device_parameters_or_buffers` function can find meta-device parameters in models and
    optimizers."""

    class BufferModule(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.register_buffer("buffer", torch.ones(2, device="meta"))

    # nn.Module
    module = torch.nn.Linear(2, 2)
    meta_module = torch.nn.Linear(2, 2, device="meta")
    buffer_meta_module = BufferModule()
    assert not _has_meta_device_parameters_or_buffers(module)
    assert _has_meta_device_parameters_or_buffers(meta_module)
    assert _has_meta_device_parameters_or_buffers(torch.nn.Sequential(module, meta_module, torch.nn.ReLU()))
    assert _has_meta_device_parameters_or_buffers(buffer_meta_module)
    # optim.Optimizer
    optimizer = torch.optim.SGD(module.parameters(), lr=0.1)
    meta_optimizer = torch.optim.SGD(meta_module.parameters(), lr=0.1)
    assert not _has_meta_device_parameters_or_buffers(optimizer)
    assert _has_meta_device_parameters_or_buffers(meta_optimizer)
    # unsupported objects
    with pytest.raises(TypeError, match="Expected `torch.nn.Module` or `torch.optim.Optimizer`"):
        _has_meta_device_parameters_or_buffers(None)
Test that the `_has_meta_device_parameters_or_buffers` function can find meta-device parameters in models and
    optimizers.
tests/tests_fabric/utilities/test_init.py