code style

colossalai/autochunk/autochunk_codegen.py

@@ -138,8 +138,8 @@ def emit_code_with_chunk(
     """
     Emit code with chunk according to chunk_infos.
 
-    It will generate a for loop in chunk regions, and replace inputs 
-        and outputs of regions with chunked variables.
+    It will generate a for loop in chunk regions, and 
+    replace inputs and outputs of regions with chunked variables.
 
     Args:
         body: forward code
@@ -157,7 +157,9 @@ def emit_code_with_chunk(
 
     # chunk inputs
     chunk_inputs = [i["inputs"] for i in chunk_infos]  # input with chunk
-    chunk_inputs_non_chunk = [i["inputs_non_chunk"] for i in chunk_infos] # input without chunk
+    chunk_inputs_non_chunk = [
+        i["inputs_non_chunk"] for i in chunk_infos
+    ]  # input without chunk
     chunk_inputs_dim = [i["inputs_dim"] for i in chunk_infos]  # input chunk dim
     chunk_inputs_names = [j.name for i in chunk_inputs for j in i] + [
         j.name for i in chunk_inputs_non_chunk for j in i