| 63 | @pytest.mark.parametrize("device", get_available_devices()) |
| 64 | @pytest.mark.parametrize("threshold", [0.0, 3.0], ids=id_formatter("threshold")) |
| 65 | def test_linear8bitlt_inference(device, threshold): |
| 66 | l1 = bnb.nn.Linear8bitLt(32, 64, threshold=threshold, has_fp16_weights=False).to(device).half() |
| 67 | assert l1.weight.device.type == device |
| 68 | assert l1.weight.dtype == torch.int8 |
| 69 | |
| 70 | l1.eval() |
| 71 | for i in range(100): |
| 72 | b1 = torch.randn(16, 8, 32, device=device).half() |
| 73 | o1 = l1(b1) |
| 74 | if i == 1: |
| 75 | assert l1.state.CB is not None |
| 76 | |
| 77 | |
| 78 | @pytest.mark.parametrize("device", get_available_devices()) |