Change the GPT-J example configure setting
@@ -555,9 +555,10 @@ def main():
return eval_loss
if model_args.tune:
- from neural_compressor.config import PostTrainingQuantConfig
+ from neural_compressor.config import AccuracyCriterion, PostTrainingQuantConfig
from neural_compressor import quantization
- conf = PostTrainingQuantConfig()
+ accuracy_criterion = AccuracyCriterion(higher_is_better=False, tolerable_loss=0.5)
+ conf = PostTrainingQuantConfig(accuracy_criterion=accuracy_criterion)
q_model = quantization.fit(model,
conf,
calib_dataloader=trainer.get_eval_dataloader(),
