Skip to content

Commit 277ac27

Browse files
authored
Correct type annotation for quantize_4bit (#994)
1 parent 8ddfda1 commit 277ac27

File tree

1 file changed

+10
-1
lines changed

1 file changed

+10
-1
lines changed

bitsandbytes/functional.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -890,7 +890,16 @@ def quantize_fp4(A: Tensor, absmax: Tensor = None, out: Tensor = None, blocksize
890890
def quantize_nf4(A: Tensor, absmax: Tensor = None, out: Tensor = None, blocksize=64, compress_statistics=False, quant_storage=torch.uint8):
891891
return quantize_4bit(A, absmax, out, blocksize, compress_statistics, 'nf4', quant_storage)
892892

893-
def quantize_4bit(A: Tensor, absmax: Tensor = None, out: Tensor = None, blocksize=64, compress_statistics=False, quant_type='fp4', quant_storage=torch.uint8) -> (Tensor, QuantState):
893+
894+
def quantize_4bit(
895+
A: Tensor,
896+
absmax: Tensor = None,
897+
out: Tensor = None,
898+
blocksize=64,
899+
compress_statistics=False,
900+
quant_type='fp4',
901+
quant_storage=torch.uint8,
902+
) -> Tuple[Tensor, QuantState]:
894903
"""
895904
Quantize tensor A in blocks of 4-bit values.
896905

0 commit comments

Comments
 (0)