File tree Expand file tree Collapse file tree 1 file changed +2
-2
lines changed
src/llmcompressor/entrypoints Expand file tree Collapse file tree 1 file changed +2
-2
lines changed Original file line number Diff line number Diff line change 4
4
5
5
import torch
6
6
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
8
8
from loguru import logger
9
9
from torch .utils .data import DataLoader
10
10
from transformers import PreTrainedModel
@@ -129,7 +129,7 @@ def __init__(
129
129
# offload to cpu if possible
130
130
if "cuda" in str (model_args .oneshot_device ) and torch .cuda .is_available ():
131
131
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 )
133
133
else :
134
134
logger .warning ("CUDA is not available! Compressing model on CPU instead" )
135
135
You can’t perform that action at this time.
0 commit comments