Python PyTorch PackageExporter.register_extern_hook用法及代码示例

本文简要介绍python语言中 torch.package.PackageExporter.register_extern_hook 的用法。

用法:

register_extern_hook(hook)

返回：

一个句柄，可用于通过调用 handle.remove() 来删除添加的钩子。

返回类型：

torch.utils.hooks.RemovableHandle

在导出器上注册一个外部挂钩。

每次模块与 extern() 模式匹配时都会调用该钩子。它应该具有以下签名：

hook(exporter: PackageExporter, module_name: str) -> None

Hooks 将按照注册的顺序被调用。