Method test_get_embedding

tests/test_adapters.py:55–61

(self)

        assert out.shape == (2, 256)

    def test_get_embedding(self):
        model = TinyModel(vocab_size=256)
        adapter = get_adapter(model)
        emb = adapter.get_embedding(model)
        ids = torch.tensor([[1, 2, 3]])
        out = emb(ids)
        assert out.shape == (1, 3, 64)

    def test_get_router_input_dim(self):
        model = TinyModel()