Remove hack comments

computational-cell-analytics · Dec 21, 2024 · 63314bd · 63314bd
1 parent 598ee6f
commit 63314bd
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/finetuning/livecell_finetuning.py b/finetuning/livecell_finetuning.py
@@ -77,7 +77,7 @@ def finetune_livecell(args):
         save_root=args.save_root,
         scheduler_kwargs=scheduler_kwargs,
         save_every_kth_epoch=args.save_every_kth_epoch,
-        peft_kwargs={"rank": args.lora_rank} if args.lora_rank is not None else None,
+        peft_kwargs={"rank": args.lora_rank, "quantize": True} if args.lora_rank is not None else None,
         mixed_precision=False,
     )
 

diff --git a/micro_sam/models/peft_sam.py b/micro_sam/models/peft_sam.py
@@ -22,7 +22,7 @@ class LoRASurgery(nn.Module):
 
     Args:
         rank: The rank of the decomposition matrices for updating weights in each attention layer.
-        block: The chosen attention blocks for implementing lora.
+        block: The chosen attention blocks for implementing LoRA.
     """
     def __init__(self, rank: int, block: nn.Module):
         super().__init__()
@@ -50,8 +50,6 @@ def forward(self, x):
         qkv = self.qkv_proj(x)  # B, N, N, 3 * org_C
         new_q = self.alpha * self.w_b_linear_q(self.w_a_linear_q(x))
         new_v = self.alpha * self.w_b_linear_v(self.w_a_linear_v(x))
-
-        # HACK
         qkv = torch.cat(
             [
                 qkv[:, :, :, :self.dim] + new_q,  # replacing new q values