Source code for ls_mlkit.util.hook.model_hook

from enum import Enum

from .base_hook import Hook, HookHandler, HookManager




[docs]
class ModelHookStageType(Enum):
    PRE_COMPUTE_LOSS = "pre_compute_loss"
    POST_COMPUTE_LOSS = "post_compute_loss"






[docs]
class ModelHookHandler(HookHandler[ModelHookStageType]):
    pass






[docs]
class ModelHook(Hook[ModelHookStageType]):
    pass






[docs]
class ModelHookManager(HookManager[ModelHookStageType]):
    pass