From 1951f7fa87725b6cc719226d26e5734958adffac Mon Sep 17 00:00:00 2001 From: oahzxl Date: Mon, 9 Jan 2023 16:30:16 +0800 Subject: [PATCH] code style --- colossalai/autochunk/autochunk_codegen.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/colossalai/autochunk/autochunk_codegen.py b/colossalai/autochunk/autochunk_codegen.py index 5ef560ac2..cc39e391e 100644 --- a/colossalai/autochunk/autochunk_codegen.py +++ b/colossalai/autochunk/autochunk_codegen.py @@ -137,9 +137,9 @@ 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 @@ -156,9 +156,11 @@ def emit_code_with_chunk( chunk_ends = [i["region"][1] for i in chunk_infos] # 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_dim = [i["inputs_dim"] for i in chunk_infos] # input chunk dim + 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_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 ]