ColossalAI/colossalai/engine/ophooks/__init__.py

4 lines
137 B
Python
Raw Normal View History

2022-04-11 15:13:02 +00:00
from .utils import register_ophooks_recursively, BaseOpHook
2022-04-11 15:13:02 +00:00
__all__ = ["BaseOpHook", "MemTracerOpHook", "register_ophooks_recursively"]