vendored model

sharp/utils/module_surgery.py

@@ -0,0 +1,32 @@
+"""Contains utility functionality to modify torch modules.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from torch import nn
+
+NORM_LAYER_TYPES = tuple(module_type for name, module_type in nn.__dict__.items() if "Norm" in name)
+BATCH_NORM_LAYER_TYPES = tuple(
+    module_type for name, module_type in nn.__dict__.items() if "BatchNorm" in name
+)
+
+
+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