diff options
Diffstat (limited to 'bitsandbytes/__init__.py')
-rw-r--r-- | bitsandbytes/__init__.py | 11 |
1 files changed, 8 insertions, 3 deletions
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 - } + } |