Function freeze_norm_layer

src/sharp/utils/module_surgery.py:19–32

Freeze all normalization layers.

(module: nn.Module)

def freeze_norm_layer(module: nn.Module) -> nn.Module:
    """Freeze all normalization layers."""

    def set_module_eval_mode(module: nn.Module, _: Any) -> None:
        module.eval()

    for submodule in module.modules():
        if isinstance(submodule, NORM_LAYER_TYPES):
            submodule.requires_grad_(False)
            # This is to ensure that batch norm layers are always called
            # with the precomputed running statistics.
            submodule.register_forward_pre_hook(set_module_eval_mode)

    return module