def test_embedder(layer, size):
    """Tests that the embedder ``predict_step`` correctly returns the output from the requested layer."""
    model = EmbedderTestModel(
        nn.Sequential(
            nn.Linear(10, 20),
            nn.Linear(20, 30),
            nn.Linear(30, 40),
        )
    )

    embedder = Embedder(model, layer)

    assert embedder.predict_step(torch.rand(10, 10), 0, 0).size(1) == size
    assert embedder(torch.rand(10, 10)).size(1) == size
Tests that the embedder ``predict_step`` correctly returns the output from the requested layer.
tests/core/utilities/test_embedder.py