Fixed issue where the CUDA SETUP was not printed.

stale
Tim Dettmers 2023-01-04 03:50:53 +07:00
parent 3901ebf7ae
commit 1341fb44ad
3 changed files with 24 additions and 1 deletions

@ -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:

@ -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!')

@ -18,7 +18,7 @@ def read(fname):
setup(
name=f"bitsandbytes",
version=f"0.36.0",
version=f"0.36.0-2",
author="Tim Dettmers",
author_email="dettmers@cs.washington.edu",
description="8-bit optimizers and matrix multiplication routines.",