ColossalAI/colossalai/engine/ophooks/__init__.py

5 lines
184 B
Python
Raw Normal View History

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