ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update utils.py

Browse source
This commit is contained in:
Daniel Han 2025-02-02 16:31:28 -08:00
commit 13c6ef7904
1 changed files with 1 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

3
unsloth/kernels/utils.py
View file

@ -154,10 +154,9 @@ 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")
ABSMAX_BUFFER = torch.empty(n_elements_absmax, dtype = torch.float32, device = "cuda:0")
if size > WEIGHT_BUFFER.numel(): WEIGHT_BUFFER.resize_(size)
if n_elements_absmax > ABSMAX_BUFFER.numel(): ABSMAX_BUFFER.resize_(n_elements_absmax)