From 66a7bd170bd8d8b590d48d4f8eca473e2a896f90 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Fri, 31 Jan 2025 09:22:14 +0100 Subject: [PATCH] style --- optimum_benchmark/backends/pytorch/backend.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/optimum_benchmark/backends/pytorch/backend.py b/optimum_benchmark/backends/pytorch/backend.py index 5b245e81..51070d8b 100644 --- a/optimum_benchmark/backends/pytorch/backend.py +++ b/optimum_benchmark/backends/pytorch/backend.py @@ -404,9 +404,9 @@ def forward(self, inputs: Dict[str, Any], kwargs: Dict[str, Any]) -> OrderedDict @torch.inference_mode() def prefill(self, inputs: Dict[str, Any], kwargs: Dict[str, Any]) -> OrderedDict: - assert ( - kwargs.get("max_new_tokens") == kwargs.get("min_new_tokens") == 1 - ), "For prefilling, max_new_tokens and min_new_tokens must be equal to 1" + assert kwargs.get("max_new_tokens") == kwargs.get("min_new_tokens") == 1, ( + "For prefilling, max_new_tokens and min_new_tokens must be equal to 1" + ) return self.pretrained_model.generate(**inputs, **kwargs) @torch.inference_mode()