github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-22 01:48:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

[zero] refactor memstats_collector (#746)

Browse Source
This commit is contained in:
HELSON
2022-04-14 12:01:12 +08:00
committed by GitHub
parent b8899e0905
commit 84c6700b2a
10 changed files with 179 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
colossalai/engine/ophooks/__init__.py
View File

@@ -1,3 +1,4 @@
from .utils import register_ophooks_recursively, BaseOpHook
from ._memtracer_ophook import MemTracerOpHook
__all__ = ["BaseOpHook", "MemTracerOpHook", "register_ophooks_recursively"]