Raise assertion error for missing installation

computational-cell-analytics · Dec 22, 2024 · bb98c16 · bb98c16
1 parent ce1f82a
commit bb98c16
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/micro_sam/models/peft_sam.py b/micro_sam/models/peft_sam.py
@@ -6,6 +6,11 @@
 
 from segment_anything.modeling import Sam
 
+try:
+    import bitsandbytes as bnb
+except Exception:
+    bnb = None
+
 
 class LoRASurgery(nn.Module):
     """Operates on the attention layers for performing low-rank adaptation.
@@ -325,6 +330,7 @@ def __init__(
         # Whether to quantize the linear layers to 4 bit precision.
         # NOTE: This is currently supported for CUDA-supported devices only.
         if quantize:
+            assert bnb is not None, "Please install 'bitsandbytes'."
             import bitsandbytes as bnb
             for name, module in model.image_encoder.named_modules():
                 if isinstance(module, torch.nn.Linear):