From 6a905be5ced93c46e35b675fbdc73d40bb95d3ee Mon Sep 17 00:00:00 2001 From: Tim Dettmers Date: Sun, 9 Jul 2023 15:34:02 -0700 Subject: [PATCH] Fixed a bug where gemv_4bit would return a wrongly sized tensor. --- bitsandbytes/functional.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/bitsandbytes/functional.py b/bitsandbytes/functional.py index aa18925..78b5f4b 100644 --- a/bitsandbytes/functional.py +++ b/bitsandbytes/functional.py @@ -1475,7 +1475,10 @@ def gemv_4bit( absmax += offset if out is None: - out = torch.zeros(size=(A.shape[0], bout), dtype=A.dtype, device=A.device) + if len(A.shape) == 3: + out = torch.zeros(size=(A.shape[0], A.shape[1], bout), dtype=A.dtype, device=A.device) + else: + out = torch.zeros(size=(A.shape[0], bout), dtype=A.dtype, device=A.device)