From 91c1964ba404dff7d3427341e92b3ca4ac26d575 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 2 Feb 2025 16:29:26 -0800 Subject: [PATCH] Update utils.py --- unsloth/kernels/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index d470c0f87e..c378d4d731 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -154,12 +154,13 @@ if HAS_CUDA_STREAM: size = shape[0]*shape[1] global WEIGHT_BUFFER global ABSMAX_BUFFER + print(size, shape) if WEIGHT_BUFFER is None: WEIGHT_BUFFER = torch.empty(size, dtype = dtype, device = "cuda:0") ABSMAX_BUFFER = torch.empty(n_elements_absmax, dtype = dtype, device = "cuda:0") if size > WEIGHT_BUFFER.numel(): WEIGHT_BUFFER.resize_(size) - if n_elements_absmax > ABSMAX_BUFFER.numel(): WEIGHT_BUFFER.resize_(n_elements_absmax) + if n_elements_absmax > ABSMAX_BUFFER.numel(): ABSMAX_BUFFER.resize_(n_elements_absmax) out = WEIGHT_BUFFER[:size].view(shape) out_absmax = ABSMAX_BUFFER[:n_elements_absmax] @@ -219,7 +220,6 @@ else: # Use same buffers for faster inference size = shape[0]*shape[1] - print(shape, size) global WEIGHT_BUFFER global ABSMAX_BUFFER if WEIGHT_BUFFER is None: