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