on running my model it showed me an error I think due to an issue in one of these files

nas.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import random
import numpy as np
from skopt import gp_minimize
from skopt.space import Real, Integer, Categorical
import math
import copy
from torch import Tensor
from torch.optim import Adam
from torch.utils.data import DataLoader
from xview3.centernet.bboxer.multilabel_circle_coder import MultilabelCircleNetCoder

class ProgressiveNAS:
    def __init__(self, base_model, search_space):
        self.base_model = base_model
        self.search_space = search_space
        self.box_coder = MultilabelCircleNetCoder(
            image_size=(1024, 1024),
            output_stride=2,
            max_objects=512,
            heatmap_encoding="umich",
            labels_encoding="circle",
            ignore_value=-1
        )
        self.population_size = 20
        self.tournament_size = 5

    def search(self, num_iterations, train_loader, val_loader):
        population = [self.mutate_model(self.base_model) for _ in range(self.population_size)]
        
        for _ in range(num_iterations):
            scores = [self.evaluate_model(model, train_loader, val_loader) for model in population]
            
            new_population = []
            for _ in range(self.population_size):
                parent1 = self.tournament_selection(population, scores)
                parent2 = self.tournament_selection(population, scores)
                child = self.crossover(parent1, parent2)
                child = self.mutate_model(child)
                new_population.append(child)
            
            population = new_population

        best_model = max(population, key=lambda model: self.evaluate_model(model, train_loader, val_loader))
        return best_model

    def tournament_selection(self, population, scores):
        tournament = random.sample(list(zip(population, scores)), self.tournament_size)
        return max(tournament, key=lambda x: x[1])[0]

    def crossover(self, parent1, parent2):
        child = copy.deepcopy(parent1)
        for name, param in child.named_parameters():
            if random.random() < 0.5:
                param.data = parent2.state_dict()[name].clone()
        return child

    def mutate_model(self, model):
        new_model = copy.deepcopy(model)
        for param in new_model.parameters():
            if random.random() < 0.1:
                param.data += torch.randn_like(param.data) * 0.1
        
        # Structure mutation
        if random.random() < 0.2:
            new_layer = nn.Linear(64, 64)  # Example layer, adjust as needed
            if hasattr(new_model, 'encoder') and hasattr(new_model.encoder, 'add_module'):
                new_model.encoder.add_module(f"extra_layer_{id(new_layer)}", new_layer)
            elif hasattr(new_model, 'decoder') and hasattr(new_model.decoder, 'add_module'):
                new_model.decoder.add_module(f"extra_layer_{id(new_layer)}", new_layer)
        
        return new_model

    def evaluate_model(self, model, train_loader, val_loader):
        # Set the model to evaluation mode
        model.eval()
        total_loss = 0.0
        total_correct = 0
        total_samples = 0

        with torch.no_grad():
            for batch_idx, batch in enumerate(train_loader):
                # Debug batch structure
                print(f"[DEBUG] Batch index: {batch_idx}")
                print(f"[DEBUG] Batch keys: {list(batch.keys())}")
                
                # Extract inputs
                inputs = batch.get('INPUT_IMAGE_KEY')
                if inputs is None:
                    raise ValueError("Batch does not contain 'INPUT_IMAGE_KEY'. Available keys: {}".format(list(batch.keys())))
                
                # Extract targets
                targets = {
                    'CENTERNET_TARGET_OBJECTNESS_MAP': batch.get('CENTERNET_TARGET_OBJECTNESS_MAP'),
                    'CENTERNET_TARGET_VESSEL_MAP': batch.get('CENTERNET_TARGET_VESSEL_MAP'),
                    'CENTERNET_TARGET_FISHING_MAP': batch.get('CENTERNET_TARGET_FISHING_MAP'),
                    'CENTERNET_TARGET_SIZE': batch.get('CENTERNET_TARGET_SIZE'),
                    'CENTERNET_TARGET_OFFSET': batch.get('CENTERNET_TARGET_OFFSET'),
                }
                # Ensure all targets are present
                if None in targets.values():
                    missing_keys = [k for k, v in targets.items() if v is None]
                    raise ValueError(f"Batch missing target keys: {missing_keys}")
                
                # Move data to device
                device = next(model.parameters()).device
                inputs = inputs.to(device)
                targets = {k: v.to(device) for k, v in targets.items()}

                # Proceed with evaluation
                outputs = model(inputs)
                loss = self.calculate_loss(outputs, targets)

                # Update metrics (This will depend on your task)
                total_loss += loss.item() * inputs.size(0)
                total_samples += inputs.size(0)
                # For object detection, calculating 'accuracy' might not be straightforward.
                # You may need to implement a custom metric, e.g., Average Precision.
                
        # Return evaluation score (e.g., total loss)
        average_loss = total_loss / total_samples
        return -average_loss  # Negative loss if you're maximizing the evaluation score
        
    def calculate_loss(self, outputs, targets):
        # Assuming outputs and targets are dictionaries with matching keys
        # You should define the loss functions for each output-target pair
        total_loss = 0

        # Heatmap loss (e.g., Focal Loss)
        heatmap_loss = self.focal_loss(outputs['heatmap'], targets['CENTERNET_TARGET_OBJECTNESS_MAP'])
        total_loss += heatmap_loss

        # Size loss (e.g., L1 Loss)
        size_loss = F.l1_loss(outputs['size'], targets['CENTERNET_TARGET_SIZE'])
        total_loss += size_loss

        # Offset loss (e.g., L1 Loss)
        offset_loss = F.l1_loss(outputs['offset'], targets['CENTERNET_TARGET_OFFSET'])
        total_loss += offset_loss

        return total_loss

    def focal_loss(self, pred, gt):
        # Implement Focal Loss as used in CenterNet
        pos_inds = gt.eq(1).float()
        neg_inds = gt.lt(1).float()
        neg_weights = torch.pow(1 - gt, 4)

        pred = torch.clamp(pred, 1e-6, 1 - 1e-6)
        pos_loss = -torch.log(pred) * torch.pow(1 - pred, 2) * pos_inds
        neg_loss = -torch.log(1 - pred) * torch.pow(pred, 2) * neg_weights * neg_inds

        num_pos  = pos_inds.float().sum()
        loss = (pos_loss + neg_loss).sum() / num_pos.clamp(min=1)
        return loss

class NAS(nn.Module):
    def __init__(self, search_space, evaluation_function):
        super(NAS, self).__init__()
        self.search_space = search_space
        self.evaluation_function = evaluation_function
        self.layers = nn.ModuleList()
        self.output_channels = None

        # Build the architecture based on the search space
        self.build_architecture()
    
    def build_architecture(self):
        # Sample architecture from the search space
        architecture = self.sample_architecture()

        # Extract parameters
        num_layers = int(architecture.get('num_layers', 2))
        num_filters = int(architecture.get('num_filters', 64))
        kernel_size = int(architecture.get('kernel_size', 3))
        activation = architecture.get('activation', 'relu')

        # Set input_channels (assuming neck_features has 256 channels)
        input_channels = 256

        # Build layers
        for i in range(num_layers):
            conv = nn.Conv2d(
                in_channels=input_channels,
                out_channels=num_filters,
                kernel_size=kernel_size,
                padding=kernel_size // 2,
                bias=False
            )
            bn = nn.BatchNorm2d(num_filters)
            act = nn.ReLU(inplace=True) if activation == 'relu' else nn.LeakyReLU(0.1, inplace=True)
            self.layers.append(nn.Sequential(conv, bn, act))
            input_channels = num_filters  # Update for next layer

        # Set output_channels based on the final layer
        self.output_channels = num_filters

    def sample_architecture(self):
        # For simplicity, we'll randomly sample from the search space
        architecture = {}
        for key, value_range in self.search_space.items():
            if isinstance(value_range, list):
                architecture[key] = random.choice(value_range)
            elif isinstance(value_range, tuple):
                architecture[key] = random.randint(*value_range)
        return architecture

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

    def search(self, num_iterations, train_loader, val_loader):
        population = self.initialize_population()
        
        for _ in range(num_iterations):
            scores = [self.evaluate_model(model, train_loader, val_loader) for model in population]
            best_model = population[np.argmax(scores)]
            
            new_population = [best_model]
            while len(new_population) < len(population):
                parent1, parent2 = self.selection(population, scores)
                child = self.crossover(parent1, parent2)
                child = self.mutate(child)
                new_population.append(child)
            
            population = new_population
        
        return population[np.argmax([self.evaluate_model(model, train_loader, val_loader) for model in population])]


    def apply_to_model(self, model, architecture):
        for name, module in model.named_modules():
            if isinstance(module, nn.Conv2d):
                if 'num_filters' in architecture:
                    in_channels = module.in_channels
                    out_channels = architecture['num_filters']
                    kernel_size = architecture.get('kernel_size', module.kernel_size)
                    new_conv = nn.Conv2d(in_channels, out_channels, kernel_size, padding=kernel_size[0]//2)
                    setattr(model, name, new_conv)
            elif isinstance(module, nn.ReLU):
                if architecture.get('activation') == 'leaky_relu':
                    setattr(model, name, nn.LeakyReLU())
        
        if 'num_layers' in architecture:
            current_layers = len(list(model.children()))
            if architecture['num_layers'] > current_layers:
                for _ in range(architecture['num_layers'] - current_layers):
                    model.add_module(f'extra_layer_{_}', nn.Linear(64, 64))  # Adjust as needed
            elif architecture['num_layers'] < current_layers:
                for _ in range(current_layers - architecture['num_layers']):
                    model = nn.Sequential(*list(model.children())[:-1])
        
        return model

# Example usage
base_model = nn.Sequential(
    nn.Conv2d(3, 64, 3, padding=1),
    nn.ReLU(),
    nn.Conv2d(64, 64, 3, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(2)
)

search_space = {
    'num_layers': (2, 5),
    'num_filters': (32, 128),
    'kernel_size': (3, 5),
    'activation': ['relu', 'leaky_relu']
}

nas = ProgressiveNAS(base_model, search_space)
# Ensure to provide train_loader and val_loader when calling search
# best_model = nas.search(num_iterations=10, train_loader=train_loader, val_loader=val_loader)

common.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import random
import numpy as np
from skopt import gp_minimize
from skopt.space import Real, Integer, Categorical
import math
import copy
from torch import Tensor
from torch.optim import Adam
from torch.utils.data import DataLoader
from xview3.centernet.bboxer.multilabel_circle_coder import MultilabelCircleNetCoder

class FilterResponseNormalization(nn.BatchNorm2d):
    def __init__(self, num_features, eps=1e-6, learnable_eps=False):
        super().__init__(num_features, eps=eps, affine=True, track_running_stats=False)
        self.learnable_eps = learnable_eps
        if learnable_eps:
            self.epsilon = nn.Parameter(torch.ones(1) * eps)

    def forward(self, x):
        mean = x.mean(dim=[2, 3], keepdim=True)
        var = x.var(dim=[2, 3], keepdim=True)
        x = (x - mean) / torch.sqrt(var + self.eps if not self.learnable_eps else self.epsilon)
        return self.weight.view(1, -1, 1, 1) * x + self.bias.view(1, -1, 1, 1)

class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = nn.Conv2d(in_channels, out_channels, 1) if in_channels != out_channels else None

    def forward(self, x):
        identity = x
        out = self.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        if self.downsample is not None:
            identity = self.downsample(x)
        out += identity

        return self.relu(out)

class FeaturePyramidNetwork(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, 1)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
        
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        return x

class AdaptiveVesselDetectionModule(nn.Module):
    def __init__(self, in_channels, num_classes):
        super().__init__()
        from xview3.centernet.models.centernet_unet import CenterNetUNetModel
        self.backbone = CenterNetUNetModel(in_channels, num_classes)
        self.adaptive_pool = nn.AdaptiveAvgPool2d((1, 1))
        self.classifier = nn.Linear(512, num_classes)

    def forward(self, x):
        features = self.backbone(x)
        pooled = self.adaptive_pool(features[-1])
        return self.classifier(pooled.view(pooled.size(0), -1))

class ProgressiveNAS:
    def __init__(self, base_model, search_space):
        self.base_model = base_model
        self.search_space = search_space
        self.box_coder = MultilabelCircleNetCoder(
            image_size=(1024, 1024),
            output_stride=2,
            max_objects=512,
            heatmap_encoding="umich",
            labels_encoding="circle",
            ignore_value=-1
        )
        self.population_size = 20
        self.tournament_size = 5

    def search(self, num_iterations, train_loader, val_loader):
        population = [self.mutate_model(self.base_model) for _ in range(self.population_size)]
        
        for _ in range(num_iterations):
            scores = [self.evaluate_model(model, train_loader, val_loader) for model in population]
            
            new_population = []
            for _ in range(self.population_size):
                parent1 = self.tournament_selection(population, scores)
                parent2 = self.tournament_selection(population, scores)
                child = self.crossover(parent1, parent2)
                child = self.mutate_model(child)
                new_population.append(child)
            
            population = new_population

        best_model = max(population, key=lambda model: self.evaluate_model(model, train_loader, val_loader))
        return best_model

    def tournament_selection(self, population, scores):
        tournament = random.sample(list(zip(population, scores)), self.tournament_size)
        return max(tournament, key=lambda x: x[1])[0]

    def crossover(self, parent1, parent2):
        child = copy.deepcopy(parent1)
        for name, param in child.named_parameters():
            if random.random() < 0.5:
                param.data = parent2.state_dict()[name].clone()
        return child

    def mutate_model(self, model):
        new_model = copy.deepcopy(model)
        for param in new_model.parameters():
            if random.random() < 0.1:
                param.data += torch.randn_like(param.data) * 0.1
        
        # Structure mutation
        if random.random() < 0.2:
            new_layer = nn.Linear(64, 64)  # Example layer, adjust as needed
            if hasattr(new_model, 'encoder') and hasattr(new_model.encoder, 'add_module'):
                new_model.encoder.add_module(f"extra_layer_{id(new_layer)}", new_layer)
            elif hasattr(new_model, 'decoder') and hasattr(new_model.decoder, 'add_module'):
                new_model.decoder.add_module(f"extra_layer_{id(new_layer)}", new_layer)
        
        return new_model

    def evaluate_model(self, model, train_loader, val_loader):
        # Set the model to evaluation mode
        model.eval()
        total_loss = 0.0
        total_correct = 0
        total_samples = 0

        with torch.no_grad():
            for batch_idx, batch in enumerate(train_loader):
                # Debug batch structure
                print(f"[DEBUG] Batch index: {batch_idx}")
                print(f"[DEBUG] Batch keys: {list(batch.keys())}")
                
                # Extract inputs
                inputs = batch.get('INPUT_IMAGE_KEY')
                if inputs is None:
                    raise ValueError("Batch does not contain 'INPUT_IMAGE_KEY'. Available keys: {}".format(list(batch.keys())))
                
                # Extract targets
                targets = {
                    'CENTERNET_TARGET_OBJECTNESS_MAP': batch.get('CENTERNET_TARGET_OBJECTNESS_MAP'),
                    'CENTERNET_TARGET_VESSEL_MAP': batch.get('CENTERNET_TARGET_VESSEL_MAP'),
                    'CENTERNET_TARGET_FISHING_MAP': batch.get('CENTERNET_TARGET_FISHING_MAP'),
                    'CENTERNET_TARGET_SIZE': batch.get('CENTERNET_TARGET_SIZE'),
                    'CENTERNET_TARGET_OFFSET': batch.get('CENTERNET_TARGET_OFFSET'),
                }
                # Ensure all targets are present
                if None in targets.values():
                    missing_keys = [k for k, v in targets.items() if v is None]
                    raise ValueError(f"Batch missing target keys: {missing_keys}")
                
                # Move data to device
                device = next(model.parameters()).device
                inputs = inputs.to(device)
                targets = {k: v.to(device) for k, v in targets.items()}

                # Proceed with evaluation
                outputs = model(inputs)
                loss = self.calculate_loss(outputs, targets)

                # Update metrics (This will depend on your task)
                total_loss += loss.item() * inputs.size(0)
                total_samples += inputs.size(0)
                # For object detection, calculating 'accuracy' might not be straightforward.
                # You may need to implement a custom metric, e.g., Average Precision.
                
        # Return evaluation score (e.g., total loss)
        average_loss = total_loss / total_samples
        return -average_loss  # Negative loss if you're maximizing the evaluation score
        
    def calculate_loss(self, outputs, targets):
        # Assuming outputs and targets are dictionaries with matching keys
        # You should define the loss functions for each output-target pair
        total_loss = 0

        # Heatmap loss (e.g., Focal Loss)
        heatmap_loss = self.focal_loss(outputs['heatmap'], targets['CENTERNET_TARGET_OBJECTNESS_MAP'])
        total_loss += heatmap_loss

        # Size loss (e.g., L1 Loss)
        size_loss = F.l1_loss(outputs['size'], targets['CENTERNET_TARGET_SIZE'])
        total_loss += size_loss

        # Offset loss (e.g., L1 Loss)
        offset_loss = F.l1_loss(outputs['offset'], targets['CENTERNET_TARGET_OFFSET'])
        total_loss += offset_loss

        return total_loss

    def focal_loss(self, pred, gt):
        # Implement Focal Loss as used in CenterNet
        pos_inds = gt.eq(1).float()
        neg_inds = gt.lt(1).float()
        neg_weights = torch.pow(1 - gt, 4)

        pred = torch.clamp(pred, 1e-6, 1 - 1e-6)
        pos_loss = -torch.log(pred) * torch.pow(1 - pred, 2) * pos_inds
        neg_loss = -torch.log(1 - pred) * torch.pow(pred, 2) * neg_weights * neg_inds

        num_pos  = pos_inds.float().sum()
        loss = (pos_loss + neg_loss).sum() / num_pos.clamp(min=1)
        return loss


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
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,  # Add this parameter
        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)]
        )
        upsample_block = {
            "nearest": nn.UpsamplingNearest2d,
            "bilinear": nn.UpsamplingBilinear2d,
            "rdtsc": ResidualDeconvolutionUpsample2d,
            "shuffle": nn.PixelShuffle,
        }[upsample_type]

        if hasattr(encoder, 'channels'):
            encoder_channels = list(encoder.channels)
        else:
            encoder_channels = [32, 24, 40, 112, 320]  # Default for EfficientNet
        encoder_channels += [encoder_channels[-1]] * num_extra_blocks

        self.decoder_channels = decoder_features

        self.decoder = D.UNetDecoder(
            encoder_channels,
            self.decoder_channels,
            unet_block=unet_block,
            upsample_block=upsample_block,
        )

        self.neck = FPNNeck(self.decoder_channels, neck_out_channels)
        self.head = head
        self.head_output_channels = head.out_channels  # Ensure your head module defines 'out_channels'

        # Ensure self.nas is a module that can be called during forward
        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  # Define this in your NAS module

        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=256, num_classes=num_classes)
        self.classifier_output_channels = self.classifier.out_channels
        self.box_coder = box_coder

        # Now define the fusion layer with known channel sizes
        total_input_channels = (
            self.head_output_channels +
            self.nas_output_channels +
            self.yolo_output_channels +
            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.batch_norm = FilterResponseNormalization(256)

    def evaluate_architecture(self, architecture):
        # Implement a simple evaluation function
        return random.random()  # Replace with actual evaluation logic

    def forward(self, x):
        x = self.initial_upsample(x)
        print(f"[DEBUG] x shape after initial_upsample: {x.shape}")
        feature_maps = self.encoder(x)
        print(f"[DEBUG] feature_maps type: {type(feature_maps)}")
        if isinstance(feature_maps, list):
            for idx, fmap in enumerate(feature_maps):
                print(f"[DEBUG] feature_maps[{idx}] shape: {fmap.shape}")
        else:
            print("[ERROR] feature_maps is not a list")
            print(f"feature_maps shape: {feature_maps.shape}")
            print(f"feature_maps type: {type(feature_maps)}")
            # Convert feature_maps to a list if needed
            feature_maps = [feature_maps]
        for extra_block in self.extra_encoder_blocks:
            feature_maps.append(extra_block(feature_maps[-1]))
        decoded_features = self.decoder(feature_maps)
        print(f"[DEBUG] decoded_features type: {type(decoded_features)}")
        if isinstance(decoded_features, list):
            for idx, fmap in enumerate(decoded_features):
                print(f"[DEBUG] decoded_features[{idx}] shape: {fmap.shape}")
        else:
            print("[ERROR] decoded_features is not a list")
            print(f"decoded_features shape: {decoded_features.shape}")
            print(f"decoded_features type: {type(decoded_features)}")
            # Convert decoded_features to a list if needed
            decoded_features = [decoded_features]
        neck_features = self.neck(decoded_features)
        print(f"[DEBUG] neck_features shape before batch_norm: {neck_features.shape}")
        neck_features = self.batch_norm(neck_features)
        print(f"[DEBUG] neck_features shape after batch_norm: {neck_features.shape}")
        neck_features = self.dropout(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)

        # Fusion layer
        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):
        # Extract and pop the encoder configuration
        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)

        # Instantiate the custom encoder
        encoder_module = CustomTimmEncoder(
            model_name=model_name,
            num_channels=num_channels,
            features_only=True,
            pretrained=pretrained
        )

        # Pop other configurations
        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)])

        # Pop decoder configuration
        decoder_config = config.pop('decoder')
        decoder_features = decoder_config['channels']
        block_type = decoder_config.get("block_type", "UnetBlock")
        upsample_type = decoder_config.get("upsample_type", "bilinear")

        neck_out_channels = 256  # Define neck output channels

        # Instantiate head and box_coder
        head_config = config.pop('head')
        head = instantiate(head_config, channels=neck_out_channels)  # Use 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),
        )

        # Now initialize the model instance with the remaining config
        return cls(
            encoder=encoder_module,
            decoder_features=decoder_features,
            head=head,
            box_coder=box_coder,
            neck_out_channels=neck_out_channels,  # Pass neck_out_channels here
            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  # Remaining config
        )

def adjust_encoder_input_channels(encoder_module, in_chans):
    # Access the underlying encoder model
    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
        print(f"[DEBUG] Before adjustment, encoder.conv_stem.in_channels: {old_conv.in_channels}")

        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
        )

        # Initialize new_conv weights
        with torch.no_grad():
            # Copy weights for existing channels
            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:
                # Initialize added 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)


the error itself


DEBUG:rasterio._env:GDAL_DATA found in environment.
DEBUG:rasterio._env:PROJ_DATA found in environment.
DEBUG:rasterio._env:Started GDALEnv: self=<rasterio._env.GDALEnv object at 0x31a4c4d00>.
DEBUG:rasterio.env:Entered env context: <rasterio.env.Env object at 0x31a4c4940>
DEBUG:rasterio._io:Output nodata value read from file: -32768.0
DEBUG:rasterio._io:Output nodata values: [-32768.0]
DEBUG:rasterio._io:all_valid: False
DEBUG:rasterio._io:mask_flags: ([<MaskFlags.nodata: 8>],)
DEBUG:rasterio._io:Jump straight to _read()
DEBUG:rasterio._io:Window: Window(col_off=18563, row_off=7781, width=512, height=512)
DEBUG:rasterio._io:IO window xoff=18563.0 yoff=7781.0 width=512.0 height=512.0
DEBUG:rasterio.env:Exiting env context: <rasterio.env.Env object at 0x31a4c4940>
DEBUG:rasterio.env:Cleared existing <rasterio._env.GDALEnv object at 0x31a4c4d00> options
DEBUG:rasterio._env:Stopped GDALEnv <rasterio._env.GDALEnv object at 0x31a4c4d00>.
DEBUG:rasterio.env:Exiting outermost env
DEBUG:rasterio.env:Exited env context: <rasterio.env.Env object at 0x31a4c4940>
[DEBUG] Channel 'vh' image shape before transformations: (512, 512)
[DEBUG] Channel 'vh' image shape after individual_transforms: (512, 512)
[DEBUG] Channel 'vh' image shape after normalization: (512, 512)
[DEBUG] Channel 'vv' image shape before transformations: (512, 512)
[DEBUG] Channel 'vv' image shape after individual_transforms: (512, 512)
[DEBUG] Channel 'vv' image shape after normalization: (512, 512)
[DEBUG] Number of channels: 2
[DEBUG] Channel 0 shape: (512, 512)
[DEBUG] Channel 1 shape: (512, 512)
[DEBUG] Combined image shape after np.dstack: (512, 512, 2)
[DEBUG] Keypoints being passed to albumentations:
(235.0, 253.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(401.0, 190.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(242.0, 487.0, 7.317978858947754) [<class 'float'>, <class 'float'>, <class 'float'>]
Error executing job with overrides: ['+config=configs/multilabel_circle_net.yaml']
Traceback (most recent call last):
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 529, in <module>
    main()
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/main.py", line 94, in decorated_main
    _run_hydra(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 394, in _run_hydra
    _run_app(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 457, in _run_app
    run_and_report(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 223, in run_and_report
    raise ex
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 220, in run_and_report
    return func()
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 458, in <lambda>
    lambda: hydra.run(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/hydra.py", line 132, in run
    _ = ret.return_value
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/core/utils.py", line 260, in return_value
    raise self._return_value
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/core/utils.py", line 186, in run_job
    ret.return_value = task_function(task_cfg)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 525, in main
    CenternetPipeline(cfg).train()
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 416, in train
    model = self.build_model(config, self.loaders)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 207, in build_model
    model = nas.search(config.nas.num_iterations, loaders['train'], loaders['valid'])
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/nas.py", line 34, in search
    scores = [self.evaluate_model(model, train_loader, val_loader) for model in population]
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/nas.py", line 34, in <listcomp>
    scores = [self.evaluate_model(model, train_loader, val_loader) for model in population]
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/nas.py", line 113, in evaluate_model
    outputs = model(inputs)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/centernet/models/centernet_unet.py", line 152, in forward
    neck_features = self.batch_norm(neck_features)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/common.py", line 26, in forward
    return self.weight.view(1, -1, 1, 1) * x + self.bias.view(1, -1, 1, 1)
RuntimeError: The size of tensor a (256) must match the size of tensor b (1280) at non-singleton dimension 1
(sam3) samerkharboush@Samers-MBP SAR-MPS-OptunaNasFbnYoloTransformers %