github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[NFC] global vars should be upper case (#1456)

pull/1543/head
Jiarui Fang 2022-08-15 09:50:29 +08:00 committed by GitHub
parent 367c615818
commit b1553fdf96
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
colossalai/fx/codegen/activation_checkpoint_codegen.py
View File

@ -4,13 +4,13 @@ from typing import List, Callable, Any, Tuple, Dict
try: try:
from torch.fx.node import Node, Argument, map_arg, _type_repr, _get_qualified_name from torch.fx.node import Node, Argument, map_arg, _type_repr, _get_qualified_name
from torch.fx.graph import _Namespace, PythonCode, _custom_builtins, _is_from_torch, _format_target, magic_methods, CodeGen, _origin_type_map, inplace_methods from torch.fx.graph import _Namespace, PythonCode, _custom_builtins, _is_from_torch, _format_target, magic_methods, CodeGen, _origin_type_map, inplace_methods
codegen_available = True CODEGEN_AVAILABLE = True
except: except:
from torch.fx.graph import _Namespace, PythonCode, _custom_builtins, _is_from_torch, _format_target, magic_methods, _origin_type_map, _format_args from torch.fx.graph import _Namespace, PythonCode, _custom_builtins, _is_from_torch, _format_target, magic_methods, _origin_type_map, _format_args
from torch.fx.node import Node, Argument, map_arg, _type_repr, _get_qualified_name from torch.fx.node import Node, Argument, map_arg, _type_repr, _get_qualified_name
codegen_available = False CODEGEN_AVAILABLE = False
if codegen_available: if CODEGEN_AVAILABLE:
__all__ = ['ActivationCheckpointCodeGen'] __all__ = ['ActivationCheckpointCodeGen']
else: else:
__all__ = ['python_code_with_activation_checkpoint'] __all__ = ['python_code_with_activation_checkpoint']
@ -169,7 +169,7 @@ def emit_code_with_activation_checkpoint(body, nodes, emit_node_func, delete_unu
within_ckpt_region = False within_ckpt_region = False
if codegen_available: if CODEGEN_AVAILABLE:
class ActivationCheckpointCodeGen(CodeGen): class ActivationCheckpointCodeGen(CodeGen):