Error in custom module

[someuser2026]

Hi, I am integrating Timm backbones with the current framework. However I am getting an error while using the pre built heads with it during inference.

Specifically it is about the channel mismatch between input to head and expected input. Attached are the erro logs as well the custom module definitions.

Thank you!

Error log

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[6], line 2
      1 # model = YOLO("../../configs/ultralytics/misc/test-backbone.yaml", task="detect")
----> 2 model = YOLO("../../configs/ultralytics/timm/dinov3_freeze-yolo11.yaml", task="segment")
      3 # model = YOLO("../../configs/ultralytics/convnext/base/convnext_base-pafpn-seg.yaml", task="obb")
      4 # print(model)
      5 img = torch.randn(1, 3, 640, 640).clip(0, 1)

File ~/Desktop/PhD/orig_paper_code/ultralytics/ultralytics/models/yolo/model.py:81, in YOLO.__init__(self, model, task, verbose)
     78     self.__dict__ = new_instance.__dict__
     79 else:
     80     # Continue with default YOLO initialization
---> 81     super().__init__(model=model, task=task, verbose=verbose)
     82     if hasattr(self.model, "model") and "RTDETR" in self.model.model[-1]._get_name():  # if RTDETR head
     83         from ultralytics import RTDETR

File ~/Desktop/PhD/orig_paper_code/ultralytics/ultralytics/engine/model.py:149, in Model.__init__(self, model, task, verbose)
    147 __import__("os").environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"  # to avoid deterministic warnings
    148 if str(model).endswith((".yaml", ".yml")):
--> 149     self._new(model, task=task, verbose=verbose)
    150 else:
    151     self._load(model, task=task)

File ~/Desktop/PhD/orig_paper_code/ultralytics/ultralytics/engine/model.py:261, in Model._new(self, cfg, task, model, verbose)
    259 self.cfg = cfg
    260 self.task = task or guess_model_task(cfg_dict)
--> 261 self.model = (model or self._smart_load("model"))(cfg_dict, verbose=verbose and RANK == -1)  # build model
    262 self.overrides["model"] = self.cfg
    263 self.overrides["task"] = self.task

File ~/Desktop/PhD/orig_paper_code/ultralytics/ultralytics/nn/tasks.py:581, in SegmentationModel.__init__(self, cfg, ch, nc, verbose)
    571 def __init__(self, cfg="yolo11n-seg.yaml", ch=3, nc=None, verbose=True):
    572     """
    573     Initialize Ultralytics YOLO segmentation model with given config and parameters.
    574 
   (...)    579         verbose (bool): Whether to display model information.
    580     """
--> 581     super().__init__(cfg=cfg, ch=ch, nc=nc, verbose=verbose)

File ~/Desktop/PhD/orig_paper_code/ultralytics/ultralytics/nn/tasks.py:435, in DetectionModel.__init__(self, cfg, ch, nc, verbose)
    433 self.model.eval()  # Avoid changing batch statistics until training begins
    434 m.training = True  # Setting it to True to properly return strides
--> 435 m.stride = torch.tensor([s / x.shape[-2] for x in _forward(torch.zeros(1, ch, s, s))])  # forward
    436 self.stride = m.stride
    437 self.model.train()  # Set model back to training(default) mode

File ~/Desktop/PhD/orig_paper_code/ultralytics/ultralytics/nn/tasks.py:431, in DetectionModel.__init__.<locals>._forward(x)
    429 if self.end2end:
    430     return self.forward(x)["one2many"]
--> 431 return self.forward(x)[0] if isinstance(m, (Segment, YOLOESegment, Pose, OBB)) else self.forward(x)

File ~/Desktop/PhD/orig_paper_code/ultralytics/ultralytics/nn/tasks.py:153, in BaseModel.forward(self, x, *args, **kwargs)
    151 if isinstance(x, dict):  # for cases of training and validating while training.
    152     return self.loss(x, *args, **kwargs)
--> 153 return self.predict(x, *args, **kwargs)

File ~/Desktop/PhD/orig_paper_code/ultralytics/ultralytics/nn/tasks.py:171, in BaseModel.predict(self, x, profile, visualize, augment, embed)
    169 if augment:
    170     return self._predict_augment(x)
--> 171 return self._predict_once(x, profile, visualize, embed)

File ~/Desktop/PhD/orig_paper_code/ultralytics/ultralytics/nn/tasks.py:194, in BaseModel._predict_once(self, x, profile, visualize, embed)
    192 if profile:
    193     self._profile_one_layer(m, x, dt)
--> 194 x = m(x)  # run
    195 y.append(x if m.i in self.save else None)  # save output
    196 if visualize:

File ~/miniforge3/envs/ultralytics_contrib/lib/python3.11/site-packages/torch/nn/modules/module.py:1773, in Module._wrapped_call_impl(self, *args, **kwargs)
   1771     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1772 else:
-> 1773     return self._call_impl(*args, **kwargs)

File ~/miniforge3/envs/ultralytics_contrib/lib/python3.11/site-packages/torch/nn/modules/module.py:1784, in Module._call_impl(self, *args, **kwargs)
   1779 # If we don't have any hooks, we want to skip the rest of the logic in
   1780 # this function, and just call forward.
   1781 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1782         or _global_backward_pre_hooks or _global_backward_hooks
   1783         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1784     return forward_call(*args, **kwargs)
   1786 result = None
   1787 called_always_called_hooks = set()

File ~/Desktop/PhD/orig_paper_code/ultralytics/ultralytics/nn/modules/conv.py:80, in Conv.forward(self, x)
     70 def forward(self, x):
     71     """
     72     Apply convolution, batch normalization and activation to input tensor.
     73 
   (...)     78         (torch.Tensor): Output tensor.
     79     """
---> 80     return self.act(self.bn(self.conv(x)))

File ~/miniforge3/envs/ultralytics_contrib/lib/python3.11/site-packages/torch/nn/modules/module.py:1773, in Module._wrapped_call_impl(self, *args, **kwargs)
   1771     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1772 else:
-> 1773     return self._call_impl(*args, **kwargs)

File ~/miniforge3/envs/ultralytics_contrib/lib/python3.11/site-packages/torch/nn/modules/module.py:1784, in Module._call_impl(self, *args, **kwargs)
   1779 # If we don't have any hooks, we want to skip the rest of the logic in
   1780 # this function, and just call forward.
   1781 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1782         or _global_backward_pre_hooks or _global_backward_hooks
   1783         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1784     return forward_call(*args, **kwargs)
   1786 result = None
   1787 called_always_called_hooks = set()

File ~/miniforge3/envs/ultralytics_contrib/lib/python3.11/site-packages/torch/nn/modules/conv.py:548, in Conv2d.forward(self, input)
    547 def forward(self, input: Tensor) -> Tensor:
--> 548     return self._conv_forward(input, self.weight, self.bias)

File ~/miniforge3/envs/ultralytics_contrib/lib/python3.11/site-packages/torch/nn/modules/conv.py:543, in Conv2d._conv_forward(self, input, weight, bias)
    531 if self.padding_mode != "zeros":
    532     return F.conv2d(
    533         F.pad(
    534             input, self._reversed_padding_repeated_twice, mode=self.padding_mode
   (...)    541         self.groups,
    542     )
--> 543 return F.conv2d(
    544     input, weight, bias, self.stride, self.padding, self.dilation, self.groups
    545 )

RuntimeError: Given groups=1, weight of size [32, 32, 1, 1], expected input[1, 48, 32, 32] to have 32 channels, but got 48 channels instead

Debug statements from parse_model printing "ch"

Unexpected keys (bn2.num_batches_tracked, bn2.bias, bn2.running_mean, bn2.running_var, bn2.weight, classifier.bias, classifier.weight, conv_head.weight) found while loading pretrained weights. This may be expected if model is being adapted.
------------------------------
[[24, 32, 48, 136, 384]]
------------------------------
------------------------------
[[24, 32, 48, 136, 384], 32]
------------------------------
------------------------------
[[24, 32, 48, 136, 384], 32, 48]
------------------------------
------------------------------
[[24, 32, 48, 136, 384], 32, 48, 136]
------------------------------
------------------------------
[[24, 32, 48, 136, 384], 32, 48, 136, 32]
------------------------------
------------------------------
[[24, 32, 48, 136, 384], 32, 48, 136, 32, 32]
------------------------------

Config Yaml

# number of classes
nc: 1

# backbone
backbone:
  - [-1, 1, Timm, ['efficientnet_b3', True, 3, True, [1,2,3,4], 32, null, 'auto', False, True]]

head:
  # - [0, 1, nn.Identity, []]
  - [0, 1, Index, [1]]
  - [0, 1, Index, [2]]
  - [0, 1, Index, [3]]

  - [1, 1, Conv, [32]]

  - [[4, 2, 3], 1, Segment, [nc, 32, 256]]

parse_model

The main change is in parsing of Index module and the addition of branch to parse Timm one along with the debug statements

def parse_model(d, ch, verbose=True):
    """
    Parse a YOLO model.yaml dictionary into a PyTorch model.

    Args:
        d (dict): Model dictionary.
        ch (int): Input channels.
        verbose (bool): Whether to print model details.

    Returns:
        model (torch.nn.Sequential): PyTorch model.
        save (list): Sorted list of output layers.
    """
    import ast

    # Args
    legacy = True  # backward compatibility for v3/v5/v8/v9 models
    max_channels = float("inf")
    nc, act, scales = (d.get(x) for x in ("nc", "activation", "scales"))
    depth, width, kpt_shape = (d.get(x, 1.0) for x in ("depth_multiple", "width_multiple", "kpt_shape"))
    if scales:
        scale = d.get("scale")
        if not scale:
            scale = tuple(scales.keys())[0]
            LOGGER.warning(f"no model scale passed. Assuming scale='{scale}'.")
        depth, width, max_channels = scales[scale]

    if act:
        Conv.default_act = eval(act)  # redefine default activation, i.e. Conv.default_act = torch.nn.SiLU()
        if verbose:
            LOGGER.info(f"{colorstr('activation:')} {act}")  # print

    if verbose:
        LOGGER.info(f"\n{'':>3}{'from':>20}{'n':>3}{'params':>10}  {'module':<45}{'arguments':<30}")
    ch = [ch]
    layers, save, c2 = [], [], ch[-1]  # layers, savelist, ch out
    base_modules = frozenset(
        {
            Classify,
            Conv,
            ConvTranspose,
            GhostConv,
            Bottleneck,
            GhostBottleneck,
            SPP,
            SPPF,
            C2fPSA,
            C2PSA,
            DWConv,
            Focus,
            BottleneckCSP,
            C1,
            C2,
            C2f,
            C3k2,
            RepNCSPELAN4,
            ELAN1,
            ADown,
            AConv,
            SPPELAN,
            C2fAttn,
            C3,
            C3TR,
            C3Ghost,
            torch.nn.ConvTranspose2d,
            DWConvTranspose2d,
            C3x,
            RepC3,
            PSA,
            SCDown,
            C2fCIB,
            A2C2f,
            # ConvNeXt-compatible modules
            # ConvNeXtStem,
            # ConvNeXtDownsample,
            # ConvNeXtBlock,
        }
    )
    repeat_modules = frozenset(  # modules with 'repeat' arguments
        {
            BottleneckCSP,
            C1,
            C2,
            C2f,
            C3k2,
            C2fAttn,
            C3,
            C3TR,
            C3Ghost,
            C3x,
            RepC3,
            C2fPSA,
            C2fCIB,
            C2PSA,
            A2C2f,
        }
    )
    for i, (f, n, m, args) in enumerate(d["backbone"] + d["head"]):  # from, number, module, args
        m = (
            getattr(torch.nn, m[3:])
            if "nn." in m
            else getattr(__import__("torchvision").ops, m[16:])
            if "torchvision.ops." in m
            else globals()[m]
        )  # get module
        for j, a in enumerate(args):
            if isinstance(a, str):
                with contextlib.suppress(ValueError):
                    args[j] = locals()[a] if a in locals() else ast.literal_eval(a)
        n = n_ = max(round(n * depth), 1) if n > 1 else n  # depth gain
        if m in base_modules:
            c1, c2 = ch[f], args[0]
            if c2 != nc:  # if c2 not equal to number of classes (i.e. for Classify() output)
                c2 = make_divisible(min(c2, max_channels) * width, 8)
            if m is C2fAttn:  # set 1) embed channels and 2) num heads
                args[1] = make_divisible(min(args[1], max_channels // 2) * width, 8)
                args[2] = int(max(round(min(args[2], max_channels // 2 // 32)) * width, 1) if args[2] > 1 else args[2])

            args = [c1, c2, *args[1:]]
            if m in repeat_modules:
                args.insert(2, n)  # number of repeats
                n = 1
            if m is C3k2:  # for M/L/X sizes
                legacy = False
                if scale in "mlx":
                    args[3] = True
            if m is A2C2f:
                legacy = False
                if scale in "lx":  # for L/X sizes
                    args.extend((True, 1.2))
            if m is C2fCIB:
                legacy = False
        elif m is AIFI:
            args = [ch[f], *args]
        elif m in frozenset({HGStem, HGBlock}):
            c1, cm, c2 = ch[f], args[0], args[1]
            args = [c1, cm, c2, *args[2:]]
            if m is HGBlock:
                args.insert(4, n)  # number of repeats
                n = 1
        elif m is ResNetLayer:
            c2 = args[1] if args[3] else args[1] * 4
        elif m is torch.nn.BatchNorm2d:
            args = [ch[f]]
        elif m is Concat:
            c2 = sum(ch[x] for x in f)
        elif m in frozenset(
            {Detect, WorldDetect, YOLOEDetect, Segment, YOLOESegment, Pose, OBB, ImagePoolingAttn, v10Detect}
        ):
            # print("f:", f)
            # print("ch:", ch)
            args.append([ch[x] for x in f])
            if m is Segment or m is YOLOESegment:
                args[2] = make_divisible(min(args[2], max_channels) * width, 8)
            if m in {Detect, YOLOEDetect, Segment, YOLOESegment, Pose, OBB}:
                m.legacy = legacy
        elif m is RTDETRDecoder:  # special case, channels arg must be passed in index 1
            args.insert(1, [ch[x] for x in f])
        elif m is CBLinear:
            c2 = args[0]
            c1 = ch[f]
            args = [c1, c2, *args[1:]]
        elif m is CBFuse:
            c2 = ch[f[-1]]
        elif m in frozenset({TorchVision}):
            c2 = args[0]
            c1 = ch[f]
            args = [*args[1:]]
        elif m in frozenset({Index}):
            c1 = ch[f]
            idx = args[0]
            if isinstance(c1, list):  # Handle multi-scale from Timm
                c2 = c1[idx]
            else:
                c2 = c1
            args = [idx]
        elif m in frozenset({ConvNeXtDownsample, ConvNeXtStem}):
            c1, c2 = ch[f], args[0]
            args = [c1, c2]
            # print("-"*50)
            # print("ConvNeXtDownsample/Stem", c1, c2, args)
            # print("-"*50)
        elif m in frozenset({ConvNeXtBlock}):
            c1, c2 = ch[f], args[0]
            drop_path_method = args[5] if len(args) > 5 else "constant"  # Fixed: should be args[5] not args[4]
    
            if n > 1 and drop_path_method == "linear":
                # Calculate progressive drop_path values
                start_prob = args[2]  # Initial drop_path value from YAML
                end_prob = args[4]    # max_drop_path from YAML
                drop_probs = [start_prob + (end_prob - start_prob) * i / (n - 1) for i in range(n)]
                
                # Create individual blocks with different drop_path values
                blocks = []
                for j in range(n):  # Fixed: use j instead of i to avoid conflict
                    # ConvNeXtBlock constructor: c1, c2, layer_scale_init_value, drop_path, use_grn, max_drop_path, drop_path_method
                    block_args = [c1, c2, args[1], drop_probs[j], args[3], args[4], args[5]]
                    # print("-"*50)
                    # print("Block args:", block_args)
                    # print("-"*50)
                    blocks.append(m(*block_args))
                m_ = torch.nn.Sequential(*blocks)
                # n = 1
            else:
                # Standard handling for constant drop_path
                args = [c1, c2, *args[1:]]
                m_ = torch.nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)
        elif m is Timm:
            c1 = ch[f]
            args[2] = c1

            m_ = m(*args)  # ← Instantiate here!
    
            # Extract channels immediately after instantiation
            if hasattr(m_, 'channels') and isinstance(m_.channels, list):
                c2 = m_.channels  # [128, 256, 512, 1024]
            else:
                raise ValueError("Timm module did not provide channel info")
        else:
            c2 = ch[f]

        # Fixed: Move this outside ConvNeXtBlock handling and fix the condition
        if m not in frozenset({ConvNeXtBlock, Timm}):
            m_ = torch.nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)  # module
        
        t = str(m)[8:-2].replace("__main__.", "")  # module type
        m_.np = sum(x.numel() for x in m_.parameters())  # number params
        m_.i, m_.f, m_.type = i, f, t  # attach index, 'from' index, type
        if verbose:
            LOGGER.info(f"{i:>3}{str(f):>20}{n_:>3}{m_.np:10.0f}  {t:<45}{str(args):<30}")  # print
        save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1)  # append to savelist
        layers.append(m_)
        if i == 0:
            ch = []
        ch.append(c2)
        print("-"*30, flush=True)
        print(ch, flush=True)
        print("-"*30)
    return torch.nn.Sequential(*layers), sorted(save)

Custom Timm module

class Timm(nn.Module):
    """
    Timm module to allow loading any timm model as a feature extractor.

    This class provides a way to load a model from the timm library with pretrained weights,
    customize input channels with intelligent weight adaptation, and extract multi-scale 
    features for tasks like detection and segmentation.

    When features_only=True, returns a list of feature tensors at different scales.
    Use with Index module to select specific scales in YAML configurations.

    Attributes:
        m (nn.Module): The loaded timm model configured as a feature extractor.
        out_indices (tuple): Indices of stages to extract features from.
        channels (list): Output channel dimensions for each feature stage.
        strides (list): Spatial reduction factors for each feature stage.

    Args:
        model (str): Name of the timm model to load (e.g., 'convnext_base', 'efficientnet_b0').
        pretrained (bool, optional): Whether to load pretrained weights. Default is True.
        in_chans (int, optional): Number of input channels. Default is 3.
        features_only (bool, optional): If True, extracts multi-scale features as a list. Default is True.
        out_indices (tuple, optional): Indices of feature stages to output. Default is (1, 2, 3, 4).
        output_stride (int, optional): Output stride for the model. Default is 32.
        norm_layer (str, optional): Normalization layer to use. Default is None (model default).
        stem_mode (str, optional): How to handle non-3 channel inputs when pretrained=True:
            - 'auto': Automatically adjust stem - repeat for >3 channels, mean for <3 (default)
            - 'repeat': Repeat RGB weights across channels
            - 'mean': Average RGB weights across channels
            - 'kaiming': Reinitialize stem with Kaiming initialization
            Note: Ignored if in_chans=3 or pretrained=False
        freeze_stem (bool, optional): Whether to freeze stem weights. Default is False.
        freeze (bool, optional): Whether to freeze entire backbone. Default is False.

    Example:
        YAML usage with Index to select specific feature scales:
        ```yaml
        backbone:
          - [-1, 1, Timm, ['convnext_base', True, 3, True, [1,2,3,4]]]  # Layer 0: [P2, P3, P4, P5]
        
        head:
          - [0, 1, Index, [3]]  # Layer 1: Select P5 (32x downsample)
          - [0, 1, Index, [2]]  # Layer 2: Select P4 (16x downsample)
          - [[1, 2], 1, Concat, [1]]  # Layer 3: Concatenate P5 and P4
        ```
    """

    def __init__(
        self,
        model: str,
        pretrained: bool = True,
        in_chans: int = 3,
        features_only: bool = True,
        out_indices: tuple = (1, 2, 3, 4),
        output_stride: int = 32,
        norm_layer: Optional[str] = None,
        stem_mode: str = 'auto',
        freeze_stem: bool = False,
        freeze: bool = False,
    ):
        """
        Load the model from timm with specified configuration.

        Args:
            model (str): Name of the timm model to load.
            pretrained (bool): Whether to load pretrained weights.
            in_chans (int): Number of input channels (RGB=3, RGBD=4, grayscale=1, etc.).
            features_only (bool): Whether to extract multi-scale features.
            out_indices (tuple): Which feature stages to output (typically 1-4 for P2-P5).
            output_stride (int): Output stride for the model.
            norm_layer (str): Normalization layer override.
            stem_mode (str): Method for adapting pretrained weights to non-3 channel inputs.
            freeze_stem (bool): Whether to freeze stem weights.
            freeze (bool): Whether to freeze all backbone parameters.
        """
        import timm  # scope for faster 'import ultralytics'

        super().__init__()
        
        # Handle non-standard input channels with pretrained weights
        if pretrained and in_chans != 3:
            # Load with 3 channels first to get pretrained weights
            model_kwargs = {
                'pretrained': True,
                'in_chans': 3,
                'features_only': features_only,
            }
            
            if features_only:
                model_kwargs['out_indices'] = out_indices
                model_kwargs['output_stride'] = output_stride
            
            if norm_layer is not None:
                model_kwargs['norm_layer'] = norm_layer
            
            self.m = timm.create_model(model, **model_kwargs)
            
            # Adapt the stem for different input channels
            self._adapt_stem(in_chans, stem_mode)
        else:
            # Standard loading (no stem adaptation needed)
            model_kwargs = {
                'pretrained': pretrained,
                'in_chans': in_chans,
                'features_only': features_only,
            }
            
            if features_only:
                model_kwargs['out_indices'] = out_indices
                model_kwargs['output_stride'] = output_stride
            
            if norm_layer is not None:
                model_kwargs['norm_layer'] = norm_layer
            
            self.m = timm.create_model(model, **model_kwargs)
        
        # Apply freezing
        if freeze:
            self._freeze_backbone()
        elif freeze_stem:
            self._freeze_stem()
        
        # Store configuration
        self.features_only = features_only
        self.out_indices = out_indices
        
        # Get feature info if available
        if features_only and hasattr(self.m, 'feature_info'):
            self.feature_info = self.m.feature_info
            self.channels = [info['num_chs'] for info in self.feature_info]
            self.strides = [info['reduction'] for info in self.feature_info]
        else:
            self.feature_info = None
            self.channels = None
            self.strides = None
    
    def _adapt_stem(self, in_chans: int, stem_mode: str):
        """
        Adapt the stem layer for different input channels while preserving pretrained weights.
        
        Args:
            in_chans (int): Target number of input channels.
            stem_mode (str): Method for weight adaptation.
        """
        # Find the first conv layer
        first_conv = None
        for name, module in self.m.named_modules():
            if isinstance(module, nn.Conv2d):
                first_conv = (name, module)
                break
        
        if first_conv is None:
            raise ValueError("Could not find first conv layer in model")
        
        name, conv = first_conv
        old_weight = conv.weight.data
        
        # Create new conv layer with target input channels
        new_conv = nn.Conv2d(
            in_chans,
            conv.out_channels,
            kernel_size=conv.kernel_size,
            stride=conv.stride,
            padding=conv.padding,
            bias=conv.bias is not None
        )
        
        # Handle weight initialization based on stem_mode
        with torch.no_grad():
            if stem_mode == 'repeat' or (stem_mode == 'auto' and in_chans > 3):
                # Repeat weights across channels (good for RGBD, multispectral, etc.)
                repeat_times = (in_chans + 2) // 3
                repeated = old_weight.repeat(1, repeat_times, 1, 1)
                new_conv.weight.data = repeated[:, :in_chans, :, :]
            elif stem_mode == 'mean' or (stem_mode == 'auto' and in_chans < 3):
                # Average weights (good for grayscale)
                new_conv.weight.data = old_weight.mean(dim=1, keepdim=True).repeat(1, in_chans, 1, 1)
            elif stem_mode == 'kaiming':
                # Reinitialize from scratch
                nn.init.kaiming_normal_(new_conv.weight, mode='fan_out', nonlinearity='relu')
            
            if conv.bias is not None:
                new_conv.bias.data = conv.bias.data
        
        # Replace the conv layer in the model
        parent_name = '.'.join(name.split('.')[:-1]) if '.' in name else ''
        attr_name = name.split('.')[-1]
        
        if parent_name:
            parent = dict(self.m.named_modules())[parent_name]
        else:
            parent = self.m
        
        setattr(parent, attr_name, new_conv)
    
    def _freeze_stem(self):
        """Freeze the stem/first few layers only."""
        frozen = False
        for name, param in self.m.named_parameters():
            if 'stem' in name.lower() or 'conv1' in name.lower() or name.startswith('0.'):
                param.requires_grad = False
                frozen = True
            elif frozen:
                break
    
    def _freeze_backbone(self):
        """Freeze all parameters in the backbone."""
        for param in self.m.parameters():
            param.requires_grad = False
        # Set model to eval mode to freeze batch norm statistics
        self.m.eval()
    
    def train(self, mode: bool = True):
        """
        Override train mode to keep backbone in eval if frozen.
        
        Args:
            mode (bool): Whether to set training mode (True) or evaluation mode (False).
        
        Returns:
            self
        """
        super().train(mode)
        # Keep backbone in eval mode if any parameters are frozen
        if any(not p.requires_grad for p in self.m.parameters()):
            self.m.eval()
        return self

    def forward(self, x: torch.Tensor) -> Union[torch.Tensor, List[torch.Tensor]]:
        """
        Forward pass through the model.

        Args:
            x (torch.Tensor): Input tensor of shape (B, C, H, W).

        Returns:
            (torch.Tensor | List[torch.Tensor]): 
                - If features_only=True: List of feature tensors at different scales
                - If features_only=False: Single output tensor
        """
        return self.m(x)
    
    def get_channel_info(self) -> Optional[dict]:
        """
        Get information about output channels and strides.

        Returns:
            (dict | None): Dictionary containing 'channels' and 'strides' lists, or None if not available.
        """
        if self.channels is not None and self.strides is not None:
            return {
                'channels': self.channels,
                'strides': self.strides,
                'out_indices': self.out_indices
            }
        return None

[UltralyticsAssistant]

👋 Hello @someuser2026, thank you for your interest in Ultralytics 🚀! This is an automated response to help you get unstuck quickly — an Ultralytics engineer will also assist you here soon.

We recommend a visit to the Docs for new users where you can find many Python and CLI usage examples and where many of the most common questions may already be answered.

If this is a 🐛 Bug Report, please provide a minimum reproducible example to help us debug it.

If this is a custom training ❓ Question, please provide as much information as possible, including dataset image examples and training logs, and verify you are following our Tips for Best Training Results.

Join the Ultralytics community where it suits you best. For real-time chat, head to Discord 🎧. Prefer in-depth discussions? Check out Discourse. Or dive into threads on our Subreddit to share knowledge with the community.

Upgrade

Upgrade to the latest ultralytics package including all requirements in a Python>=3.8 environment with PyTorch>=1.8 to verify your issue is not already resolved in the latest version:

pip install -U ultralytics

Environments

YOLO may be run in any of the following up-to-date verified environments (with all dependencies including CUDA/CUDNN, Python and PyTorch preinstalled):

• Notebooks with free GPU:
• Google Cloud Deep Learning VM. See GCP Quickstart Guide
• Amazon Deep Learning AMI. See AWS Quickstart Guide
• Docker Image. See Docker Quickstart Guide

Status

If this badge is green, all Ultralytics CI tests are currently passing. CI tests verify correct operation of all YOLO Modes and Tasks on macOS, Windows, and Ubuntu every 24 hours and on every commit.

[someuser2026]

solved

[glenn-jocher]

Glad to hear it’s resolved—please add a brief note on the root cause/fix and mark this discussion as Answered; if your parse_model() or Timm/Index handling is generally useful, a small, focused PR with tests and docs would be welcome.

[someuser2026]

The timm module was storing info about all the channels in the self.channels list instead of for only the ones that are specified in the out_indices.