torch.nn.modules.module.register_module_buffer_registration_hook

torch.nn.modules.module.register_module_buffer_registration_hook(hook)[源代码]

注册一个适用于所有模块的缓冲区注册钩子。

警告

这向nn.Module模块添加了全局状态。

每次调用 register_buffer() 时都会触发此钩子。它应具有以下签名:

hook(module, name, buffer) -> None or new buffer

该钩子可以修改输入,并返回一个修改后的值。

返回值

一个可以通过调用 handle.remove() 移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

本页目录