From fd818cf14423489714cd3fb19c703a5b40271e17 Mon Sep 17 00:00:00 2001 From: oahzxl Date: Tue, 10 Jan 2023 11:10:45 +0800 Subject: [PATCH] change imports --- colossalai/autochunk/autochunk_codegen.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/colossalai/autochunk/autochunk_codegen.py b/colossalai/autochunk/autochunk_codegen.py index 1ee1d818a..14f17b1d3 100644 --- a/colossalai/autochunk/autochunk_codegen.py +++ b/colossalai/autochunk/autochunk_codegen.py @@ -18,7 +18,8 @@ if CODEGEN_AVAILABLE: inplace_methods, magic_methods, ) - from torch.fx.node import Argument, Node, _get_qualified_name, _type_repr, map_arg + +from torch.fx.node import Argument, Node, _get_qualified_name, _type_repr, map_arg from .search_chunk import SearchChunk from .utils import delete_free_var_from_last_use, find_idx_by_name, get_node_shape