diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py
index 882b329d4..7874eef35 100644
--- a/codeflash/optimization/optimizer.py
+++ b/codeflash/optimization/optimizer.py
@@ -4,7 +4,9 @@ from typing import List, Tuple, Union
 import openai
 
 
-def optimize_python_code(source_code: str, n: int = 1) -> List[Tuple[Union[str, None], Union[str, None]]]:
+def optimize_python_code(
+    source_code: str, n: int = 1
+) -> List[Tuple[Union[str, None], Union[str, None]]]:
     """
     Optimize the given python code for performance using OpenAI's GPT-4 model.
 
@@ -17,13 +19,15 @@ def optimize_python_code(source_code: str, n: int = 1) -> List[Tuple[Union[str,
     """
     system_message = {
         "role": "system",
-        "content": ("You are a computer programmer who writes really fast programs and is an expert in optimizing "
-                    "runtime and memory requirements of a program by rewriting them. You don't rename function names "
-                    "or change function signatures. The function return value should be exactly the same as before.")
+        "content": (
+            "You are a computer programmer who writes really fast programs and is an expert in optimizing "
+            "runtime and memory requirements of a program by rewriting them. You don't rename function names "
+            "or change function signatures. The function return value should be exactly the same as before."
+        ),
     }
     user_message = {
         "role": "user",
-        "content": f"Rewrite this python program to run faster.\n```python\n{source_code}\n```"
+        "content": f"Rewrite this python program to run faster.\n```python\n{source_code}\n```",
     }
     messages = [system_message, user_message]
 
@@ -49,7 +53,9 @@ def optimize_python_code(source_code: str, n: int = 1) -> List[Tuple[Union[str,
 
 
 if __name__ == "__main__":
-    optimize_python_code("""
+    print(
+        optimize_python_code(
+            """
     def sorter(arr):
     n = len(arr) 
     for i in range(n): 
@@ -61,4 +67,7 @@ if __name__ == "__main__":
         if swapped == False:
             break
     return arr
-    """, 3)
+    """,
+            3,
+        )
+    )