diff --git a/src/pruna/algorithms/quanto.py b/src/pruna/algorithms/quanto.py index a5b31c56..1b90031f 100644 --- a/src/pruna/algorithms/quanto.py +++ b/src/pruna/algorithms/quanto.py @@ -44,7 +44,7 @@ class Quanto(PrunaAlgorithmBase): algorithm_name: str = "quanto" group_tags: list[tags] = [tags.QUANTIZER] references: dict[str, str] = {"GitHub": "https://github.com/huggingface/optimum-quanto"} - save_fn: SAVE_FUNCTIONS = SAVE_FUNCTIONS.reapply + save_fn: SAVE_FUNCTIONS = SAVE_FUNCTIONS.save_before_apply tokenizer_required: bool = False processor_required: bool = False dataset_required: bool = False diff --git a/src/pruna/algorithms/torch_dynamic.py b/src/pruna/algorithms/torch_dynamic.py index 77c8f625..d09923c4 100644 --- a/src/pruna/algorithms/torch_dynamic.py +++ b/src/pruna/algorithms/torch_dynamic.py @@ -37,7 +37,7 @@ class TorchDynamic(PrunaAlgorithmBase): algorithm_name = "torch_dynamic" group_tags: list[tags] = [tags.QUANTIZER] references: dict[str, str] = {"GitHub": "https://github.com/pytorch/pytorch"} - save_fn: SAVE_FUNCTIONS = SAVE_FUNCTIONS.pickled + save_fn: SAVE_FUNCTIONS = SAVE_FUNCTIONS.save_before_apply tokenizer_required: bool = False processor_required: bool = False runs_on: list[str] = ["cpu", "cuda"] diff --git a/tests/algorithms/testers/torch_dynamic.py b/tests/algorithms/testers/torch_dynamic.py index f1b1614a..a0c9c4da 100644 --- a/tests/algorithms/testers/torch_dynamic.py +++ b/tests/algorithms/testers/torch_dynamic.py @@ -6,7 +6,7 @@ class TestTorchDynamic(AlgorithmTesterBase): """Test the torch dynamic quantizer.""" - models = ["shufflenet"] + models = ["shufflenet", "smollm_135m"] reject_models = [] allow_pickle_files = False algorithm_class = TorchDynamic