huggingface/pytorch-image-models · error · AssertionError
You have provided a batch norm layer as the `root module`. P
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.
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_moduleView on GitHub (pinned to 9a5261e31b)
Solutions
- Use timm.utils.model.freeze_batch_norm_2d(bn_module) or unfreeze_batch_norm_2d(bn_module) for a single BN layer
- 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)
Defensive patterns
Strategy: type-guard
Type guard
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')
Prevention
- Branch on module type before freezing
- Use freeze_(model, submodules=[...]) for bulk BN freezing
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
- Input image must have positive dimensions, got H={height}, W
- Invalid class map file, expected a dict ({class_map_path}).
- Dataset length is unknown, please pass `num_samples` explici
- Found 0 images in subfolders of {root}. Supported image exte
- Invalid or corrupt tar info cache file {cache_path}.
AI-assisted analysis of huggingface/pytorch-image-models@9a5261e31b (2026-08-27).
Data as JSON: /api/errors/b94cd1766f06a066.
Report an issue: GitHub.