Add changes suggested by linter

microsoft · Mar 20, 2024 · f0ebf4d · f0ebf4d
1 parent c44005b
commit f0ebf4d
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/onnxruntime/python/tools/transformers/models/llama/benchmark_e2e.py b/onnxruntime/python/tools/transformers/models/llama/benchmark_e2e.py
@@ -383,7 +383,11 @@ def main():
                 inputs["attention_mask"] = torch.cat(
                     [inputs["attention_mask"], (~has_eos).to(torch.int64).reshape(batch_size, 1)], 1
                 )
-                inputs["position_ids"] = None if "position_ids" not in inputs else torch.max(inputs["position_ids"], dim=1)[0].reshape(batch_size, 1) + 1
+                inputs["position_ids"] = (
+                    None
+                    if "position_ids" not in inputs
+                    else torch.max(inputs["position_ids"], dim=1)[0].reshape(batch_size, 1) + 1
+                )
 
                 # Set logits to zeros for next inference run and re-use memory buffer
                 if outputs["logits"].shape[1] != 1: