Fixed a bug where gemv_4bit would return a wrongly sized tensor.

6a905be5 · Tim Dettmers · 0f0390ac · 6a905be5
Commit 6a905be5 authored Jul 09, 2023 by Tim Dettmers
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

bitsandbytes/functional.py bitsandbytes/functional.py +4 -1

No files found.
--- 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)