check_bnb_install.py 333 Bytes
Newer Older
1
2
import torch

Aarni Koskela's avatar
Aarni Koskela committed
3
4
import bitsandbytes as bnb

Ruff's avatar
Ruff committed
5
6
p = torch.nn.Parameter(torch.rand(10, 10).cuda())
a = torch.rand(10, 10).cuda()
7
8
9
10
11

p1 = p.data.sum().item()

adam = bnb.optim.Adam([p])

Ruff's avatar
Ruff committed
12
out = a * p
13
14
15
16
17
18
19
loss = out.sum()
loss.backward()
adam.step()

p2 = p.data.sum().item()

assert p1 != p2
Ruff's avatar
Ruff committed
20
21
print("SUCCESS!")
print("Installation was successful!")