From 7f95372c6267d3163fd2aa74aeff9d84ddb6cc35 Mon Sep 17 00:00:00 2001
From: Victor Agostinelli <35327788+agostinv@users.noreply.github.com>
Date: Fri, 6 Dec 2024 04:39:45 -0800
Subject: [PATCH] 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 <93391238+MekkCyber@users.noreply.github.com>
---
 src/transformers/integrations/bitnet.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/transformers/integrations/bitnet.py b/src/transformers/integrations/bitnet.py
index 3386bdcb43b27c..0b50f9738afb69 100644
--- 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(