fix import issue

bitsandbytes-foundation · Feb 8, 2024 · 0905ad7 · 0905ad7
1 parent 8b4baaa
commit 0905ad7
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 6 deletions.
diff --git a/bitsandbytes/__init__.py b/bitsandbytes/__init__.py
@@ -17,7 +17,9 @@
 
 if COMPILED_WITH_CUDA:
     from .optim import adam
-
+    from .backends import register_backend, backends
+    from .backends.cuda import CUDABackend
+    register_backend("cuda", CUDABackend())
 __pdoc__ = {
     "libbitsandbytes": False,
     "optim.optimizer.Optimizer8bit": False,

diff --git a/bitsandbytes/backends/__init__.py b/bitsandbytes/backends/__init__.py
@@ -1,14 +1,9 @@
 from typing import Dict
 import torch
 
-from bitsandbytes.cextension import COMPILED_WITH_CUDA
 from bitsandbytes.backends.base import Backend
 
 backends: Dict[str, Backend] = {}
 
 def register_backend(backend_name: str, backend_instance: Backend):
     backends[backend_name.lower()] = backend_instance
-
-if COMPILED_WITH_CUDA:
-    from .cuda import CUDABackend
-    register_backend("cuda", CUDABackend())