From 8258b4364a21a4da2572cb644d0926080c3268da Mon Sep 17 00:00:00 2001 From: Max Ryabinin Date: Fri, 1 Jul 2022 17:16:10 +0300 Subject: Add a CPU-only build option --- bitsandbytes/__init__.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) (limited to 'bitsandbytes/__init__.py') diff --git a/bitsandbytes/__init__.py b/bitsandbytes/__init__.py index 6e29322..22fb841 100644 --- a/bitsandbytes/__init__.py +++ b/bitsandbytes/__init__.py @@ -2,9 +2,14 @@ # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. -from .optim import adam + from .nn import modules -__pdoc__ = {'libBitsNBytes' : False, +from cextension import COMPILED_WITH_CUDA + +if COMPILED_WITH_CUDA: + from .optim import adam + +__pdoc__ = {'libBitsNBytes': False, 'optim.optimizer.Optimizer8bit': False, 'optim.optimizer.MockArgs': False - } + } -- cgit v1.2.3