(device, caplog)
| 415 | |
| 416 | @pytest.mark.parametrize("device", get_available_devices()) |
| 417 | def test_4bit_linear_warnings(device, caplog): |
| 418 | dim1 = 64 |
| 419 | |
| 420 | with caplog_at_level(caplog, logging.WARNING, "bitsandbytes.nn.modules"): |
| 421 | net = nn.Sequential(*[bnb.nn.Linear4bit(dim1, dim1, quant_type="nf4") for i in range(10)]) |
| 422 | net = net.to(device) |
| 423 | inp = torch.rand(10, dim1, device=device, dtype=torch.float16) |
| 424 | net(inp) |
| 425 | assert any("inference or training" in msg for msg in caplog.messages) |
| 426 | |
| 427 | caplog.clear() |
| 428 | with caplog_at_level(caplog, logging.WARNING, "bitsandbytes.nn.modules"): |
| 429 | net = nn.Sequential(*[bnb.nn.Linear4bit(dim1, dim1, quant_type="nf4") for i in range(10)]) |
| 430 | net = net.to(device) |
| 431 | inp = torch.rand(1, dim1, device=device, dtype=torch.float16) |
| 432 | net(inp) |
| 433 | assert any("inference." in msg for msg in caplog.messages) |
| 434 | |
| 435 | |
| 436 | @pytest.mark.parametrize("device", get_available_devices()) |
nothing calls this directly
no test coverage detected