Files
ColossalAI/colossalai/engine/ophooks/__init__.py
2022-04-14 12:01:12 +08:00

5 lines
184 B
Python

from .utils import register_ophooks_recursively, BaseOpHook
from ._memtracer_ophook import MemTracerOpHook
__all__ = ["BaseOpHook", "MemTracerOpHook", "register_ophooks_recursively"]