add weight property to exllamav2 quanlinear (#80)

predibase · Nov 29, 2023 · 360ad4c · 360ad4c
1 parent cb96f12
commit 360ad4c
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/server/lorax_server/utils/gptq/exllamav2.py b/server/lorax_server/utils/gptq/exllamav2.py
@@ -136,6 +136,10 @@ def temp_fwd_size(self, max_input_len, max_batch_size):
     def scratch_space_fixed(self, max_input_len=2048, max_batch_size=8):
         return self.temp_dq_size() + self.temp_fwd_size(max_input_len, max_batch_size)
 
+    @property
+    def weight(self) -> torch.Tensor:
+        return self.qweight
+
 
 class ExLlamaV2DeviceTensors:
 
@@ -158,4 +162,4 @@ def get_scratch_slice(self, size_bytes):
         size_bytes = ((size_bytes + 127) // 128) * 128
         size_half = size_bytes // 2
         scratch_slice = self.scratch.narrow(0, 0, size_half)
-        return scratch_slice
+        return scratch_slice