Fixed issue where the CUDA SETUP was not printed.
This commit is contained in:
parent
3901ebf7ae
commit
1341fb44ad
|
@ -1,4 +1,5 @@
|
|||
import ctypes as ct
|
||||
import os
|
||||
import torch
|
||||
|
||||
from pathlib import Path
|
||||
|
@ -10,6 +11,8 @@ from bitsandbytes.cuda_setup.main import CUDASetup
|
|||
setup = CUDASetup.get_instance()
|
||||
if setup.initialized != True:
|
||||
setup.run_cuda_setup()
|
||||
if 'BITSANDBYTES_NOWELCOME' not in os.environ or str(os.environ['BITSANDBYTES_NOWELCOME']) == '0':
|
||||
setup.print_log_stack()
|
||||
|
||||
lib = setup.lib
|
||||
try:
|
||||
|
|
20
check_bnb_install.py
Normal file
20
check_bnb_install.py
Normal file
|
@ -0,0 +1,20 @@
|
|||
import bitsandbytes as bnb
|
||||
import torch
|
||||
|
||||
p = torch.nn.Parameter(torch.rand(10,10).cuda())
|
||||
a = torch.rand(10,10).cuda()
|
||||
|
||||
p1 = p.data.sum().item()
|
||||
|
||||
adam = bnb.optim.Adam([p])
|
||||
|
||||
out = a*p
|
||||
loss = out.sum()
|
||||
loss.backward()
|
||||
adam.step()
|
||||
|
||||
p2 = p.data.sum().item()
|
||||
|
||||
assert p1 != p2
|
||||
print('SUCCESS!')
|
||||
print('Installation was successful!')
|
Loading…
Reference in New Issue
Block a user