Skip to content

Commit

Permalink
fix format
Browse files Browse the repository at this point in the history
  • Loading branch information
echarlaix committed Sep 26, 2023
1 parent f7141a6 commit 72f59e2
Show file tree
Hide file tree
Showing 2 changed files with 2 additions and 2 deletions.
2 changes: 1 addition & 1 deletion tests/neural_compressor/test_optimization.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,7 @@ def test_dynamic_quantization(self, task, model_name, expected_quantized_matmuls
tokenizer = AutoTokenizer.from_pretrained(model_name)
save_onnx_model = False
quantized_model = None
model_kwargs = {"use_cache" : False, "use_io_binding": False} if task == "text-generation" else {}
model_kwargs = {"use_cache": False, "use_io_binding": False} if task == "text-generation" else {}
with tempfile.TemporaryDirectory() as tmp_dir:
for backend in ["torch", "ort"]:
if backend == "torch":
Expand Down
2 changes: 1 addition & 1 deletion tests/neural_compressor/utils_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,7 @@ def check_model_outputs(
file_name=None,
):
tokens = tokenizer("This is a sample input", return_tensors="pt")
file_name = ONNX_WEIGHTS_NAME if task!="text-generation" else "decoder_model.onnx"
file_name = ONNX_WEIGHTS_NAME if task != "text-generation" else "decoder_model.onnx"

model_kwargs = (
{"decoder_file_name": file_name, "use_cache": False, "use_io_binding": False}
Expand Down

0 comments on commit 72f59e2

Please sign in to comment.