Skip to content
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions tests/e2e/e2e_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
from llmcompressor.modifiers.quantization import GPTQModifier, QuantizationModifier
from tests.test_timer.timer_utils import log_time
from tests.testing_utils import process_dataset
from transformers import DefaultDataCollator


def load_model(model: str, model_class: str, device_map: str | None = None):
Expand Down Expand Up @@ -94,7 +95,11 @@ def data_collator(batch):
)

# Apply quantization.

logger.info("ONESHOT KWARGS", oneshot_kwargs)

oneshot_kwargs["shuffle_calibration_samples"] = True
oneshot_kwargs["data_collator"] = DefaultDataCollator()
_run_oneshot(**oneshot_kwargs)

return oneshot_kwargs["model"], processor
Loading