xnn.common.models.registry.replace_module#

xnn.common.models.registry.replace_module(chain, old, new)[source]#

Replace old by new inside a chain of wrappers linked by .model.

Parameters:

chain (torch.nn.Module) – The outermost module (old itself, or a wrapper around it).
old (torch.nn.Module) – The module to replace and its replacement.
new (torch.nn.Module) – The module to replace and its replacement.

Returns:

new when chain is old, otherwise chain with the replacement made in place.

Return type:

torch.nn.Module

Raises:

ValueError – If old is not in the chain.