Function test_4bit_embedding_warnings

tests/test_modules.py:437–448

(device, caplog)

@pytest.mark.parametrize("device", get_available_devices())
def test_4bit_embedding_warnings(device, caplog):
    num_embeddings = 128
    default_block_size = 64

    with caplog_at_level(caplog, logging.WARNING, "bitsandbytes.nn.modules"):
        net = bnb.nn.Embedding4bit(
            num_embeddings=num_embeddings, embedding_dim=default_block_size + 1, quant_type="nf4"
        )
        net.to(device)
        inp = torch.randint(low=0, high=num_embeddings, size=(1,), device=device)
        net(inp)
    assert any("inference" in msg for msg in caplog.messages)


@pytest.mark.parametrize("device", get_available_devices(no_cpu=True))