Skip to content

Commit 819df1c

Browse files
committed
rename function
Signed-off-by: Kyle Sayers <[email protected]>
1 parent ad71c5b commit 819df1c

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

src/llmcompressor/entrypoints/oneshot.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55
import torch
66
from accelerate.hooks import remove_hook_from_module
7-
from compressed_tensors.utils import force_cpu_offload
7+
from compressed_tensors.utils import offloaded_dispatch
88
from loguru import logger
99
from torch.utils.data import DataLoader
1010
from transformers import PreTrainedModel
@@ -129,7 +129,7 @@ def __init__(
129129
# offload to cpu if possible
130130
if "cuda" in str(model_args.oneshot_device) and torch.cuda.is_available():
131131
remove_hook_from_module(model_args.model, recurse=True)
132-
force_cpu_offload(model_args.model, model_args.oneshot_device)
132+
offloaded_dispatch(model_args.model, model_args.oneshot_device)
133133
else:
134134
logger.warning("CUDA is not available! Compressing model on CPU instead")
135135

0 commit comments

Comments
 (0)