nntoolbox.hooks.hooks module

Implement abstraction for hooks Adopt from FastAI:

class nntoolbox.hooks.hooks.Hook(module: torch.nn.modules.module.Module, hook_func: Callable[[nntoolbox.hooks.hooks.Hook, torch.nn.modules.module.Module, torch.Tensor, torch.Tensor], Any], forward: bool = True)[source]

Bases: object

remove()[source]
class nntoolbox.hooks.hooks.Hooks(ms: List[torch.nn.modules.module.Module], hook_fn: Callable[[nntoolbox.hooks.hooks.Hook, torch.nn.modules.module.Module, torch.Tensor, torch.Tensor], Any], forward)[source]

Bases: object

remove()[source]