diff --git a/bitsandbytes/cextension.py b/bitsandbytes/cextension.py index 00ee587..7a62c1e 100644 --- a/bitsandbytes/cextension.py +++ b/bitsandbytes/cextension.py @@ -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: diff --git a/check_bnb_install.py b/check_bnb_install.py new file mode 100644 index 0000000..77cd03e --- /dev/null +++ b/check_bnb_install.py @@ -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!') diff --git a/setup.py b/setup.py index aa01d79..93df40e 100644 --- a/setup.py +++ b/setup.py @@ -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.",