Add feature dim attributes to BitLinear for easier PEFT integration (#…

…34946) Update bitnet.py, extremely small change to allow for easier PEFT integration Co-authored-by: Mohamed Mekkouri <[email protected]>
huggingface · Dec 6, 2024 · 7f95372 · 7f95372
1 parent 9ad4c93
commit 7f95372
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/transformers/integrations/bitnet.py b/src/transformers/integrations/bitnet.py
@@ -127,6 +127,8 @@ class BitLinear(nn.Module):
     def __init__(self, in_features: int, out_features: int, bias: bool, device=None, dtype=None):
         super().__init__()
         self.dtype = dtype
+        self.in_features = in_features
+        self.out_features = out_features
         self.register_buffer(
             "weight",
             torch.zeros(