return F.conv2d(input, weight, bias, self.stride,
RuntimeError: Given groups=1, weight of size [256, 1280, 1, 1], expected input[8, 256, 32, 32] to have 1280 channels, but got 256 channels instead

I've this error and I noticed that for this [256, 1280, 1, 1],  when I change this  "neck_out_channels = config.get('neck_out_channels', 256)"  which in def from_config in the centernet_unet.py 
the 256 to any number the first number in this  [256, 1280, 1, 1] changes accordingly and the next number be 5th folds of the first number!! 
now help me  fix the mismatch between my size and expected !

here's full centernet_unet.py

from functools import partial
from typing import List, Dict
from omegaconf import DictConfig
from pytorch_toolbelt.modules import (
    ACT_RELU,
    decoders as D,
    EncoderModule,
    ResidualDeconvolutionUpsample2d,
)
from torch import nn
import torch
import timm
from hydra.utils import instantiate
import random
import torch.nn.functional as F
from xview3.centernet.models.custom_encoder import CustomTimmEncoder, DebugUNetDecoder
from .unet_blocks import IRBlock, get_unet_block, SqueezeExcitation
from .necks import FPNNeck
from xview3.nas import NAS as NeuralArchitectureSearch
from .heads.yolo_head import YOLOHead
from .transformer_classifier import ConvolutionalClassifier
from xview3.common import FilterResponseNormalization, ResidualBlock, FeaturePyramidNetwork
from ..constants import (
    CENTERNET_OUTPUT_CLASS_MAP,
    CENTERNET_OUTPUT_OBJECTNESS_MAP,
    CENTERNET_OUTPUT_OFFSET,
    CENTERNET_OUTPUT_SIZE,
    CENTERNET_OUTPUT_VESSEL_MAP,
    CENTERNET_OUTPUT_FISHING_MAP,
)
__all__ = ["CenterNetUNetModel"]

class CenterNetUNetModel(nn.Module):
    def __init__(
        self,
        encoder: EncoderModule,
        decoder_features: List[int],
        head: nn.Module,
        box_coder,
        neck_out_channels: int = 256,  
        input_upsample_factor: int = 1,
        activation=ACT_RELU,
        num_extra_blocks=0,
        block_type="UnetBlock",
        upsample_type="bilinear",
        extra_block=partial(IRBlock, stride=2),
        num_classes=2,
        anchors=[(10, 13), (16, 30), (33, 23)],
        **kwargs
    ):
        super().__init__()
        self.output_stride = getattr(encoder, 'strides', [32])[0]
        self.initial_upsample = (
            nn.UpsamplingBilinear2d(scale_factor=input_upsample_factor)
            if input_upsample_factor != 1 else nn.Identity()
        )
        unet_block = get_unet_block(block_type, activation=activation)
        self.encoder = encoder
        self.extra_encoder_blocks = nn.ModuleList(
            [extra_block(encoder.channels[-1], encoder.channels[-1]) for _ in range(num_extra_blocks)]
        )
        if upsample_type in ['nearest', 'bilinear']:
            upsample_block_cls = partial(
                nn.Upsample,
                scale_factor=2,
                mode=upsample_type,
                align_corners=False
            )
        elif upsample_type == 'rdtsc':
            upsample_block_cls = partial(ResidualDeconvolutionUpsample2d, scale_factor=2)
        elif upsample_type == 'shuffle':
            upsample_block_cls = partial(nn.PixelShuffle, upscale_factor=2)
        else:
            raise ValueError(f"Unknown upsample_type {upsample_type}")
        if hasattr(encoder, 'channels'):
            encoder_channels = list(encoder.channels)
        else:
            encoder_channels = [32, 24, 40, 112, 320]  
        encoder_channels += [encoder_channels[-1]] * num_extra_blocks

        self.head = head
        self.head_output_channels = head.out_channels 

        
        self.nas = NeuralArchitectureSearch(
            search_space={
                'num_layers': (2, 5),
                'num_filters': (32, 128),
                'kernel_size': (3, 5),
                'activation': ['relu', 'leaky_relu']
            },
            evaluation_function=self.evaluate_architecture
        )
        self.nas_output_channels = self.nas.output_channels  

        self.yolo_head = YOLOHead(num_classes=num_classes, anchors=anchors)
        self.yolo_output_channels = self.yolo_head.num_anchors * (num_classes + 5)  # Adjust as per YOLO implementation

        self.classifier = ConvolutionalClassifier(in_channels=neck_out_channels, num_classes=num_classes)
        self.classifier_output_channels = self.classifier.out_channels
        self.box_coder = box_coder
        self.head_output_channels = getattr(self.head, 'out_channels', 256)
        self.nas_output_channels = getattr(self.nas, 'output_channels', 256)
        self.yolo_output_channels = self.yolo_head.num_anchors * (num_classes + 5)
        self.classifier_output_channels = getattr(self.classifier, 'out_channels', num_classes)

       
        total_input_channels = (
            int(self.head_output_channels) +
            int(self.nas_output_channels) +
            int(self.yolo_output_channels) +
            int(self.classifier_output_channels)
        )
        self.fusion = nn.Conv2d(total_input_channels, self.head_output_channels, kernel_size=1)

        self.dropout = nn.Dropout(p=0.5)

        self.decoder_channels = decoder_features

        encoder_channels = list(encoder.channels)
        num_blocks = len(encoder_channels) - 1 

       
        if len(self.decoder_channels) < num_blocks:
           
            self.decoder_channels.extend(
                [self.decoder_channels[-1]] * (num_blocks - len(self.decoder_channels))
            )
        elif len(self.decoder_channels) > num_blocks:
            
            self.decoder_channels = self.decoder_channels[:num_blocks]

        self.num_decoded_features = len(encoder_channels)
        num_decoder_stages = self.num_decoded_features

        self.neck_output_channels = neck_out_channels * self.num_decoded_features
   
        self.decoder = DebugUNetDecoder(
            encoder_channels=encoder_channels,
            decoder_channels=self.decoder_channels,
            unet_block=unet_block,
            upsample_block=upsample_block_cls,
            center=False,
        )

        self.neck = FPNNeck(channels=self.decoder_channels, out_channels=neck_out_channels)
        self.neck_conv = nn.Conv2d(
            in_channels=self.neck_output_channels,
            out_channels=neck_out_channels,
            kernel_size=1
        )
        self.batch_norm = FilterResponseNormalization(self.neck_output_channels)

    def evaluate_architecture(self, architecture):
        
        return random.random() 

    def forward(self, x):
        x = self.initial_upsample(x)
        feature_maps = self.encoder(x)
        decoded_features = self.decoder(feature_maps)

        print(f"Decoded features length: {len(decoded_features)}")
        for idx, fmap in enumerate(decoded_features):
            print(f"[DEBUG] decoded_features[{idx}] shape: {fmap.shape}")

        neck_features = self.neck(decoded_features)

        neck_features = self.batch_norm(neck_features)
        neck_features = self.dropout(neck_features)
        neck_features = self.neck_conv(neck_features)

        head_output = self.head(neck_features)
        nas_output = self.nas(neck_features)
        yolo_output = self.yolo_head(neck_features)
        classification_output = self.classifier(neck_features)

        fused_output = self.fusion(torch.cat([head_output, nas_output, yolo_output, classification_output], dim=1))

        return {
            CENTERNET_OUTPUT_OBJECTNESS_MAP: fused_output[:, :1],
            CENTERNET_OUTPUT_VESSEL_MAP: fused_output[:, 1:2],
            CENTERNET_OUTPUT_FISHING_MAP: fused_output[:, 2:3],
            CENTERNET_OUTPUT_SIZE: fused_output[:, 3:4],
            CENTERNET_OUTPUT_OFFSET: fused_output[:, 4:6],
        }

    @classmethod
    def from_config(cls, **config):
        
        encoder_config = config.pop('encoder')
        num_channels = config.pop('num_channels', 3)
        model_name = encoder_config.get('model_name', 'tf_efficientnet_b5_ns')
        pretrained = encoder_config.get('pretrained', False)

        encoder_module = CustomTimmEncoder(
            model_name=model_name,
            num_channels=num_channels,
            features_only=True,
            pretrained=pretrained
        )

        encoder_module = adjust_encoder_input_channels(encoder_module, num_channels)

        input_upsample_factor = config.pop("input_upsample_factor", 1)
        activation = config.pop('activation', ACT_RELU)
        num_extra_blocks = config.pop('num_extra_blocks', 0)
        num_classes = config.pop("num_classes", 2)
        anchors = config.pop("anchors", [(10, 13), (16, 30), (33, 23)])

        decoder_config = config.pop('decoder')
        encoder_channels = list(encoder_module.channels)  
        decoder_features = decoder_config['channels']
        block_type = decoder_config.get("block_type", "UnetBlock")
        upsample_type = decoder_config.get("upsample_type", "bilinear")

        neck_out_channels = config.get('neck_out_channels', 256)
        neck_output_channels = neck_out_channels * len(decoder_features)
        num_blocks = len(encoder_channels) - 1

        num_decoder_stages = len(encoder_channels) - 1

        if len(decoder_features) < num_decoder_stages:
            decoder_features.extend([decoder_features[-1]] * (num_decoder_stages - len(decoder_features)))
        elif len(decoder_features) > num_decoder_stages:
            decoder_features = decoder_features[:num_decoder_stages]

        head_config = config.pop('head')
        head = instantiate(head_config, channels=neck_output_channels)

        box_coder_config = config.pop('box_coder')
        box_coder = instantiate(
            box_coder_config,
            output_stride=encoder_module.strides[0] // (head.upsample_factor * input_upsample_factor),
        )

        return cls(
            encoder=encoder_module,
            decoder_features=decoder_features,
            head=head,
            box_coder=box_coder,
            neck_out_channels=neck_out_channels,
            input_upsample_factor=input_upsample_factor,
            activation=activation,
            block_type=block_type,
            upsample_type=upsample_type,
            num_extra_blocks=num_extra_blocks,
            num_classes=num_classes,
            anchors=anchors,
            **config
        )

def adjust_encoder_input_channels(encoder_module, in_chans):
  
    if hasattr(encoder_module, 'encoder'):
        base_encoder = encoder_module.encoder
    else:
        base_encoder = encoder_module

    if hasattr(base_encoder, 'conv_stem'):
        old_conv = base_encoder.conv_stem
     
        new_conv = nn.Conv2d(
            in_channels=in_chans,
            out_channels=old_conv.out_channels,
            kernel_size=old_conv.kernel_size,
            stride=old_conv.stride,
            padding=old_conv.padding,
            bias=old_conv.bias is not None
        )
        
        with torch.no_grad():
           
            num_existing_in_channels = min(old_conv.in_channels, in_chans)
            new_conv.weight[:, :num_existing_in_channels, ...] = old_conv.weight[:, :num_existing_in_channels, ...]
            if in_chans > old_conv.in_channels:
               
                torch.nn.init.kaiming_normal_(new_conv.weight[:, old_conv.in_channels:, ...])

        base_encoder.conv_stem = new_conv
        
        print(f"[DEBUG] After adjustment, encoder.conv_stem.in_channels: {base_encoder.conv_stem.in_channels}")
    else:
        raise NotImplementedError("Adjusting input channels for this encoder is not implemented.")
    return encoder_module

class EnhancedCenterNetUNet(nn.Module):
    def __init__(self, in_channels, num_classes):
        super().__init__()
        self.backbone = CenterNetUNetModel(in_channels, num_classes)
        self.fpn = FeaturePyramidNetwork([64, 128, 256, 512], 256)
        self.se_blocks = nn.ModuleList([SqueezeExcitation(c) for c in [64, 128, 256, 512]])

    def forward(self, x):
        features = self.backbone(x)
        for i, feature in enumerate(features):
            features[i] = self.se_blocks[i](feature)
        return self.fpn(features)


#here's the config file


defaults:
  - _self_
  - dataset: valid_sar_cv
  - augs: flips_light
  - sampler: balanced
  - normalization: default
  - loss: centernet/rfl_bce_mse
  - optimizer: adamw
  - scheduler: cos
  - train: default
  - runner: default
  - ema: disabled
  - torch: default

batch_size: 8
channels_last: False

local_rank: ${oc.env:LOCAL_RANK, 0}
world_size: ${oc.env:WORLD_SIZE, 0}

experiment: null
seed: 42
deterministic: False
detect_anomaly: False
find_unused: False

transfer: null
checkpoint: null

nas:
  search_space:
    num_layers: [2, 5]
    num_filters: [32, 128]
    kernel_size: [3, 5]
    activation: ['relu', 'leaky_relu']
  num_iterations: 10

model:
  _target_: xview3.centernet.models.centernet_unet.CenterNetUNetModel.from_config
  num_channels: 2
  num_classes: 2
  anchors:
    - [10, 13]
    - [16, 30]
    - [33, 23]
  input_upsample_factor: 1
  encoder:
    _target_: custom_encoder.CustomTimmEncoder
    model_name: 'tf_efficientnet_b5_ns'
    num_channels: 2
    features_only: True
    pretrained: False
  decoder:
    channels: [512, 384, 256, 128]
    block_type: UnetBlock
    upsample_type: bilinear
  head:
    _target_: xview3.centernet.models.heads.decoupled_head.DecoupledHeadGroupNormLateShuffle
    channels: 32
    classifier_dim: 64
    dropout_rate: 0.1
  box_coder:
    _target_: xview3.centernet.bboxer.MultilabelCircleNetCoder
    image_size: [768, 768]
    output_stride: 4
    max_objects: 512
    heatmap_encoding: umich
    labels_encoding: circle
    ignore_value: -1
  activation: silu
  num_extra_blocks: 0

yolo:
  enabled: true
  loss_weight: 1.0

transformer:
  enabled: true
  loss_weight: 1.0

train:
  epochs: 4
  early_stopping: 2
  train_image_size: [768, 768]
  valid_image_size: [2048, 2048]

torch:
  cudnn:
    allow_tf32: True
  cuda_allow_tf32: True
  detect_anomaly: False
  deterministic: False
  benchmark: True

hydra:
  output_subdir: null
  run:
    dir: .
  sweep:
    dir: .
    subdir: .

augs:
  spatial:
    - albumentations.HorizontalFlip
    - albumentations.VerticalFlip
  individual:
    vv0:
      _target_: xview3.augmentations.unclipped_random_brightness.UnclippedRandomBrightnessContrast
      always_apply: false
      p: 0.5

normalization:
  channels:
    vv:
      _target_: xview3.SigmoidNormalization
      midpoint: -20.0
      temperature: 0.18
    vh:
      _target_: xview3.SigmoidNormalization
      midpoint: -20.0
      temperature: 0.18
    bathymetry:
      _target_: xview3.CubicRootNormalization
    wind_direction:
      _target_: albumentations.NoOp
    wind_speed:
      _target_: albumentations.NoOp
    mask:
      _target_: albumentations.NoOp

dataset:
  data_dir: ${oc.env:XVIEW3_DIR}/full_dataset
  splitter: valid_sar_cv
  fold: 0
  num_folds: 4

# here's neck.py


import collections
from typing import List, Dict
#from pytorch_toolbelt.modules import FPNFuse
import torch
from torch.nn import functional as F
from torch import nn, Tensor

__all__ = ["TensorAtIndex", "PickTensors", "HyperColumnNeck"]


class TensorAtIndex(nn.Module):
    def __init__(self, channels: List[int], index: int = 0):
        super().__init__()
        self.index = index
        self.channels = channels[index]

    def forward(self, feature_maps: List[Tensor]) -> Tensor:
        return feature_maps[self.index]


class PickTensors(nn.Module):
    def __init__(self, channels: List[int], indexes: Dict[int, str]):
        super().__init__()
        self.indexes = indexes
        self.channels = channels[0]

    def forward(self, feature_maps: List[Tensor]) -> Dict[str, Tensor]:
        return collections.OrderedDict((name, feature_maps[index]) for index, name in self.indexes.items())


class HyperColumnNeck(nn.Module):
    def __init__(self, channels: List[int], dropout_rate=0.0):
        super().__init__()
        self.fuse = FPNFuse(align_corners=True)
        self.channels = sum(channels)
        self.dropout = nn.Dropout2d(p=dropout_rate, inplace=True)

    def forward(self, feature_maps: List[Tensor]) -> Tensor:
        return self.dropout(self.fuse(feature_maps))

class FPNFuse(nn.Module):
    def __init__(
        self,
        input_channels: List[int],
        output_channels: int,
        mode: str = 'bilinear',
        align_corners: bool = False,
    ):
        super().__init__()
        self.mode = mode
        self.align_corners = align_corners
        self.lateral_convs = nn.ModuleList([
            nn.Conv2d(in_ch, output_channels, kernel_size=1, bias=False)
            for in_ch in input_channels
        ])
        self.num_features = len(input_channels)
        self.fuse_conv = nn.Conv2d(
            output_channels * self.num_features,
            output_channels,
            kernel_size=3,
            padding=1
        )
        
    def forward(self, feature_maps: List[Tensor]):
        print(f"[DEBUG] feature_maps in FPNFuse: type={type(feature_maps)}")
        valid_feature_maps = []
        dst_size = None
        for idx, f in enumerate(feature_maps):
            if f.dim() == 4:
                print(f"[DEBUG] Valid feature_maps[{idx}] shape: {f.shape}")
                # Apply lateral convolution
                processed_feature = self.lateral_convs[idx](f)
                valid_feature_maps.append(processed_feature)
                if dst_size is None:
                    dst_size = processed_feature.size()[-2:]
            else:
                print(f"[WARNING] Skipping feature_maps[{idx}] with shape {f.shape}")
        if dst_size is None:
            raise ValueError("No valid feature maps with spatial dimensions found.")
        layers = []
        for f in valid_feature_maps:
            upsampled_feature = F.interpolate(
                f,
                size=dst_size,
                mode=self.mode,
                align_corners=self.align_corners
            )
            layers.append(upsampled_feature)
        x = torch.cat(layers, dim=1)
        x = self.fuse_conv(x)
        return x

class FPNNeck(nn.Module):
    def __init__(self, channels: List[int], out_channels: int):
        super().__init__()
        self.fpn = FPNFuse(channels, out_channels, mode='bilinear', align_corners=False)
        self.neck_output_channels = out_channels * (len(channels)+1)
        self.out_conv = nn.Conv2d(self.neck_output_channels, out_channels, kernel_size=1)

    def forward(self, feature_maps: List[Tensor]) -> Tensor:
        fpn_output = self.fpn(feature_maps)
        print(f"[DEBUG] fpn_output shape: {fpn_output.shape}")
        return self.out_conv(fpn_output)


#here's yolo.py


from torch import nn, Tensor
import torch
#from pytorch_toolbelt.modules import FPNFuse
from typing import List, Tuple, Union
from ..necks import FPNFuse

class AnchorFreeHead(nn.Module):
    def __init__(self, in_channels, num_classes):
        super().__init__()
        self.cls_conv = nn.Conv2d(in_channels, num_classes, 1)
        self.reg_conv = nn.Conv2d(in_channels, 4, 1)
        self.obj_conv = nn.Conv2d(in_channels, 1, 1)

    def forward(self, x):
        cls_pred = self.cls_conv(x)
        reg_pred = self.reg_conv(x)
        obj_pred = self.obj_conv(x)
        return torch.cat([reg_pred, obj_pred, cls_pred], dim=1)

class YOLOHead(nn.Module):
    def __init__(self, num_classes, anchors):
        super(YOLOHead, self).__init__()
        self.num_classes = num_classes
        self.anchors = anchors
        self.num_anchors = len(anchors)
        self.fpn = FPNFuse([256, 512, 1024], 256)
        self.out_channels = self.num_anchors * (5 + self.num_classes)
        self.conv = nn.Conv2d(256, self.out_channels, 1)

    def forward(self, x):
        x = self.fpn(x)
        batch_size, _, height, width = x.shape
        return self.conv(x).view(batch_size, len(self.anchors), 5 + self.num_classes, height, width).permute(0, 1, 3, 4, 2)


#here's decoupled_head.py


from collections import OrderedDict
from typing import Dict, Tuple, List

from pytorch_toolbelt.modules import ACT_RELU, get_activation_block, conv1x1
from torch import nn, Tensor

from ...constants import (
    CENTERNET_OUTPUT_VESSEL_MAP,
    CENTERNET_OUTPUT_FISHING_MAP,
    CENTERNET_OUTPUT_OBJECTNESS_MAP,
    CENTERNET_OUTPUT_SIZE,
    CENTERNET_OUTPUT_OFFSET,
)

__all__ = ["DecoupledHeadGroupNormLateShuffle"]


class FeatureRecombination(nn.Module):
    def __init__(self, channels, modulating_features):
        super().__init__()
        self.project = conv1x1(modulating_features, channels)

    def forward(self, x, modulation):
        return x * self.project(modulation).sigmoid()


def build_regression_block(
    input_channels: int,
    embedding_size: int,
    dropout_rate: float,
    num_blocks: int,
    activation: str,
    extra_final_layers: List[Tuple[str, nn.Module]],
    num_groups=16,
):
    activation_block = get_activation_block(activation)

    layers = [
        (
            "project",
            nn.Conv2d(input_channels, embedding_size, kernel_size=(1, 1), padding=0, bias=False),
        ),
    ]

    for layer_index in range(num_blocks):
        layers += [
            (
                f"conv{layer_index}",
                nn.Conv2d(embedding_size, embedding_size, kernel_size=(3, 3), padding=1, bias=False),
            ),
            (f"bn{layer_index}", nn.GroupNorm(num_groups, embedding_size)),
            (f"act{layer_index}", activation_block(inplace=True)),
        ]
    layers += [
        ("drop", nn.Dropout2d(dropout_rate, inplace=False)),
        ("conv1", nn.Conv2d(embedding_size, embedding_size, kernel_size=(1, 1), padding=0, bias=False)),
    ] + extra_final_layers

    return nn.Sequential(OrderedDict(layers))


class DecoupledHeadGroupNormLateShuffle(nn.Module):
    upsample_factor: int = 2

    def __init__(
        self,
        channels: int,
        classifier_dim: int,
        dropout_rate=0.25,
        objectness_dim: int = 32,
        size_dim: int = 32,
        offset_dim: int = 32,
        num_blocks: int = 3,
        activation: str = ACT_RELU,
        num_groups=16,
    ):
        super().__init__()
        self.objectness = build_regression_block(
            channels,
            embedding_size=objectness_dim,
            dropout_rate=dropout_rate,
            activation=activation,
            num_groups=num_groups,
            num_blocks=num_blocks,
            extra_final_layers=[
                (("shuffle"), nn.PixelShuffle(2)),
                ("final", nn.Conv2d(objectness_dim // 4, 1, kernel_size=(3, 3), padding=1)),
            ],
        )

        self.classifier = build_regression_block(
            channels,
            embedding_size=classifier_dim,
            dropout_rate=dropout_rate,
            num_groups=num_groups,
            activation=activation,
            num_blocks=num_blocks,
            extra_final_layers=[],
        )

        self.size = build_regression_block(
            channels,
            embedding_size=size_dim,
            dropout_rate=dropout_rate,
            activation=activation,
            num_groups=num_groups,
            num_blocks=num_blocks,
            extra_final_layers=[],
        )

        self.classifier_recombination = FeatureRecombination(classifier_dim, size_dim)
        self.classifier_up = nn.Sequential(
            OrderedDict(
                [
                    ("conv1", nn.Conv2d(classifier_dim, classifier_dim, kernel_size=(3, 3), padding=1)),
                    ("shuffle", nn.PixelShuffle(2)),
                    ("final", nn.Conv2d(classifier_dim // 4, 2, kernel_size=(3, 3), padding=1)),
                ]
            )
        )
        self.size_recombination = FeatureRecombination(size_dim, classifier_dim)
        self.size_up = nn.Sequential(
            OrderedDict(
                [
                    ("conv1", nn.Conv2d(size_dim, size_dim, kernel_size=(3, 3), padding=1)),
                    ("shuffle", nn.PixelShuffle(2)),
                    ("final", nn.Conv2d(size_dim // 4, 1, kernel_size=(3, 3), padding=1)),
                    ("activation", nn.Softplus()),
                ]
            )
        )

        if offset_dim:
            self.offset = build_regression_block(
                channels,
                embedding_size=offset_dim,
                dropout_rate=dropout_rate,
                activation=activation,
                num_groups=num_groups,
                num_blocks=num_blocks,
                extra_final_layers=[
                    (("shuffle"), nn.PixelShuffle(2)),
                    ("final", nn.Conv2d(offset_dim // 4, 2, kernel_size=(3, 3), padding=1)),
                ],
            )
        else:
            self.offset = None

        # Determine if offset is included
        if self.offset is not None:
            self.out_channels = 6  # Total output channels
        else:
            self.out_channels = 4  # Exclude offset channels

        self._init_weights()

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight)
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

        # self.classifier_up.final.bias.data.fill_(0)
        # self.objectness.final.bias.data.fill_(-4)

        self.objectness.final.bias.data.fill_(-5)
        self.classifier_up.final.bias.data.fill_(-1)
        self.size_up.final.bias.data.fill_(2)

    def forward(self, features: Tensor) -> Dict[str, Tensor]:
        class_features = self.classifier(features)
        length_features = self.size(features)

        new_class_features = self.classifier_recombination(class_features, length_features)
        new_length_features = self.size_recombination(length_features, class_features)

        predicted_classes = self.classifier_up(new_class_features)
        predicted_length = self.size_up(new_length_features)

        output = {
            CENTERNET_OUTPUT_OBJECTNESS_MAP: self.objectness(features),
            CENTERNET_OUTPUT_VESSEL_MAP: predicted_classes[:, 0:1],
            CENTERNET_OUTPUT_FISHING_MAP: predicted_classes[:, 1:2],
            CENTERNET_OUTPUT_SIZE: predicted_length,
        }
        if self.offset is not None:
            output[CENTERNET_OUTPUT_OFFSET] = self.offset(features).clip(0, 1)
        return output


#here's custom_encoder.py

import timm
import torch.nn as nn
import torch
import torch.nn.functional as F
from pytorch_toolbelt.modules.decoders import UNetDecoder

class CustomTimmEncoder(nn.Module):
    def __init__(self, model_name, num_channels=2, features_only=True, pretrained=False):
        super(CustomTimmEncoder, self).__init__()
        self.encoder = timm.create_model(
            model_name,
            features_only=features_only,
            in_chans=num_channels,
            pretrained=pretrained
        )
        # Extract channels and strides
        self.channels = [f['num_chs'] for f in self.encoder.feature_info]
        self.strides = [f['reduction'] for f in self.encoder.feature_info]
        self.feature_info = self.encoder.feature_info
        
    def forward(self, x):
        features = self.encoder(x)
        for idx, feature in enumerate(features):
            print(f"[DEBUG] Feature {idx} shape: {feature.shape}")
        return features
    
class DebugUNetDecoder(nn.Module):
    def __init__(self, encoder_channels, decoder_channels, unet_block, upsample_block, center=False):
        super(DebugUNetDecoder, self).__init__()
        self.center = center
        self.num_stages = len(decoder_channels)
        self.up_samples = nn.ModuleList()
        self.unet_blocks = nn.ModuleList()

        if self.center:
            self.center_block = unet_block(encoder_channels[-1], encoder_channels[-1])
        else:
            self.center_block = nn.Identity()

        for i in range(self.num_stages):
            self.up_samples.append(upsample_block())

            if i == 0:
                in_channels = encoder_channels[-1] + encoder_channels[-2]
            else:
                in_channels = decoder_channels[i - 1] + encoder_channels[-(i + 2)]
            out_channels = decoder_channels[i]

            self.unet_blocks.append(unet_block(in_channels, out_channels))

    def forward(self, features):
        x = features[-1]
        x = self.center_block(x)
        num_features = len(features)
        print(f"[DEBUG] Number of encoder features: {num_features}")

        decoded_features = []
        for i in range(self.num_stages):
            print(f"[DEBUG] Decoder stage {i}")
            x = self.up_samples[i](x)

            encoder_index = num_features - i - 2
            if encoder_index >= 0:
                encoder_feature = features[encoder_index]
                print(f"Using encoder feature at index {encoder_index}")

                if x.size()[2:] != encoder_feature.size()[2:]:
                    encoder_feature = F.interpolate(encoder_feature, size=x.size()[2:], mode='nearest')
                    print(f"encoder_feature.shape after resizing at stage {i}: {encoder_feature.shape}")

                x = torch.cat([x, encoder_feature], dim=1)
                print(f"x.shape after concat at stage {i}: {x.shape}")
            else:
                print(f"No encoder feature for decoder stage {i}, skipping concatenation")

            x = self.unet_blocks[i](x)
            decoded_features.append(x)

        return decoded_features