github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-03 18:19:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

[fx] added activation checkpoint codegen support for torch < 1.12 (#1359)

Browse Source
This commit is contained in:
Frank Lee
2022-07-25 23:35:31 +08:00
committed by GitHub
parent 4417804129
commit cd063ac37f
3 changed files with 435 additions and 192 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/fx/codegen/__init__.py
View File

@@ -1,3 +1 @@
from .activation_checkpoint_codegen import ActivationCheckpointCodeGen
__all__ = ['ActivationCheckpointCodeGen']
from .activation_checkpoint_codegen import *