Fixed bitsandbytes imports for nf4 (#72)

predibase · Nov 27, 2023 · ffc8049 · ffc8049
1 parent 8347f58
commit ffc8049
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/server/lorax_server/utils/layers.py b/server/lorax_server/utils/layers.py
@@ -11,7 +11,7 @@
 HAS_BITS_AND_BYTES = True
 try:
     import bitsandbytes as bnb
-    from bitsandbytes.nn import Int8Params
+    from bitsandbytes.nn import Int8Params, Params4bit
 
 except ImportError:
     HAS_BITS_AND_BYTES = False