huggingface/pytorch-image-models · error · AssertionError

You have provided a batch norm layer as the `root module`…

Error message

You have provided a batch norm layer as the `root module`. Please use `timm.utils.model.freeze_batch_norm_2d` or `timm.utils.model.unfreeze_batch_norm_2d` instead.

What it means

freeze_/unfreeze_ in timm.utils.model refuse to operate when the root module passed is itself a BatchNorm layer (or timm's BatchNormAct2d variants), because in-place conversion is impossible and the operation would silently do the wrong thing. AssertionError is raised with a pointer to freeze_batch_norm_2d/unfreeze_batch_norm_2d which handle single BN layers.

Solutions

1. Use timm.utils.model.freeze_batch_norm_2d(bn_module) or unfreeze_batch_norm_2d(bn_module) for a single BN layer
2. Pass the parent model (or a container submodule) to freeze_/unfreeze_ with a submodules filter to select BN layers

Example fix

# before
from timm.utils.model import freeze_
freeze_(model.bn1)  # AssertionError
# after
from timm.utils.model import freeze_batch_norm_2d
freeze_batch_norm_2d(model.bn1)

import torch.nn as nn\nfrom timm.layers.norm import BatchNormAct2d\n\ndef is_bn(m):\n    return isinstance(m, (nn.modules.batchnorm._BatchNorm, BatchNormAct2d))\n\n# route:\nfreeze_batch_norm_2d(m) if is_bn(m) else freeze_(m, 'bn')

When it happens

Trigger: Calling freeze_(model.bn1) or unfreeze_(some_bn_module) where the argument is an nn.BatchNorm2d/BatchNormAct2d/SyncBatchNorm(Sync)Act instance instead of a container model.

Common situations: User grabbed a BN submodule from named_modules() to freeze just that layer; passing model.get_submodule('bn1') instead of the parent model.

Related errors

• A progressive schedule requires at least two choices.
• adamw_lr is deprecated, use fallback_lr_scale=adamw_lr/lr…
• All scheduled batch sizes must be positive integers.
• API has changed, `state_steps` argument must contain a list…
• batch_sizes must contain at least one value.

AI-assisted analysis of huggingface/pytorch-image-models@9a5261e31b (2026-08-27). Data as JSON: /api/errors/b94cd1766f06a066. Report an issue: GitHub.

Appendix: source

Thrown at timm/utils/model.py:134

    Args:
        root_module (nn.Module, optional): Root module relative to which the `submodules` are referenced.
        submodules (list[str]): List of modules for which the parameters will be (un)frozen. They are to be provided as
            named modules relative to the root module (accessible via `root_module.named_modules()`). An empty list
            means that the whole root module will be (un)frozen. Defaults to []
        include_bn_running_stats (bool): Whether to also (un)freeze the running statistics of batch norm 2d layers.
            Defaults to `True`.
        mode (bool): Whether to freeze ("freeze") or unfreeze ("unfreeze"). Defaults to `"freeze"`.
    """
    assert mode in ["freeze", "unfreeze"], '`mode` must be one of "freeze" or "unfreeze"'

    if isinstance(root_module, (
            torch.nn.modules.batchnorm.BatchNorm2d,
            torch.nn.modules.batchnorm.SyncBatchNorm,
            BatchNormAct2d,
            SyncBatchNormAct,
    )):
        # Raise assertion here because we can't convert it in place
        raise AssertionError(
            "You have provided a batch norm layer as the `root module`. Please use "
            "`timm.utils.model.freeze_batch_norm_2d` or `timm.utils.model.unfreeze_batch_norm_2d` instead.")

    if isinstance(submodules, str):
        submodules = [submodules]

    named_modules = submodules
    submodules = [root_module.get_submodule(m) for m in submodules]

    if not len(submodules):
        named_modules, submodules = list(zip(*root_module.named_children()))

    for n, m in zip(named_modules, submodules):
        # (Un)freeze parameters
        for p in m.parameters():
            p.requires_grad = False if mode == 'freeze' else True
        if include_bn_running_stats:
            # Helper to add submodule specified as a named_module

View on GitHub (pinned to 9a5261e31b)