github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-05-10 01:17:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' of https://github.com/hpcaitech/ColossalAI into dev0116

Browse Source
This commit is contained in:
jiaruifang 2023-01-18 18:43:11 +08:00
commit 7f822a5c45
1 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
colossalai/fx/graph_module.py
View File

@ -22,7 +22,12 @@ if COLOGM:
class ColoGraphModule(GraphModule):
def __init__(self, root: Union[torch.nn.Module, Dict[str, Any]], graph: Graph, class_name: str = 'GraphModule'):
def __init__(self,
root: Union[torch.nn.Module, Dict[str, Any]],
graph: Graph,
class_name: str = 'GraphModule',
ckpt_codegen: bool = True):
if ckpt_codegen:
graph.set_codegen(ActivationCheckpointCodeGen())
super().__init__(root, graph, class_name)