ColossalAI/colossalai/gemini/ophooks/__init__.py

5 lines
184 B
Python

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