ls_mlkit.util.hook.model_hook module

class ls_mlkit.util.hook.model_hook.ModelHook(name: str, stage: HookStageType, fn: Callable[[...], Any | None], priority: int = 0, enabled: bool = True)[source]

Bases: Hook[ModelHookStageType]

class ls_mlkit.util.hook.model_hook.ModelHookHandler(manager: HookManager[HookStageType], hook: Hook[HookStageType])[source]

Bases: HookHandler[ModelHookStageType]

class ls_mlkit.util.hook.model_hook.ModelHookManager[source]

Bases: HookManager[ModelHookStageType]

class ls_mlkit.util.hook.model_hook.ModelHookStageType(*values)[source]

Bases: Enum

POST_COMPUTE_LOSS = 'post_compute_loss'
PRE_COMPUTE_LOSS = 'pre_compute_loss'