this is the format of fashionpedia dataset :
{
 "info": info,
 "categories": [category],
 "attributes": [attribute],
 "images": [image],
 "annotations": [annotation],
 "licenses": [license]
}

info{
  "year" : int,
  "version" : str,
  "description" : str,
  "contributor" : str,
  "url" : str,
  "date_created" : datetime,
}

category{
  "id" : int,
  "name" : str,
  "supercategory" : str,  # parent of this label
  "level": int,           # levels in the taxonomy
  "taxonomy_id": string,
}

attribute{
  "id" : int,
  "name" : str,
  "supercategory" : str,  # parent of this label
  "level": int,           # levels in the taxonomy
  "taxonomy_id": string,
}

image{
  "id" : int,
  "width" : int,
  "height" : int,
  "file_name" : str,
  "license" : int,
  "time_captured": string,
  "original_url": string,
  "isstatic": int, 0: the original_url is not a static url,
  "kaggle_id": str,
}

annotation{
  "id" : int,
  "image_id" : int,
  "category_id" : int,
  "attribute_ids": [int],
  "segmentation" : [polygon] or [rle]
  "bbox" : [x,y,width,height], # int
  "area" : int
  "iscrowd": int (1 or 0)
}
polygon: [x1, y1, x2, y2, ...], where x, y are the coordinates of vertices, int
rle: {"size", (height, widht), "counts": str}

license{
  "id" : int,
  "name" : str,
  "url" : str
}
resolve this error :
in train_one_epoch(model, dataloader, optimizer, scheduler, scaler, device)
    641                             target_attrs = attribute_ids[b, :n]
    642 
--> 643                             instance_attr_loss += F.binary_cross_entropy_with_logits(
    644                                 pred_attrs,
    645                                 target_attrs,

/usr/local/lib/python3.10/dist-packages/torch/nn/functional.py in binary_cross_entropy_with_logits(input, target, weight, size_average, reduce, reduction, pos_weight)
   3622 
   3623     if not (target.size() == input.size()):
-> 3624         raise ValueError(
   3625             f"Target size ({target.size()}) must be the same as input size ({input.size()})"
   3626         )

ValueError: Target size (torch.Size([5, 9])) must be the same as input size (torch.Size([5, 294]))
my guess is model is predicting in one hot format but annotations include just id of each attribute.
Give me whole code rewritten and corrected.

the original code :
import os
import json
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from transformers import AutoImageProcessor, AutoModel  # Fixed import
from torch.amp import autocast, GradScaler
from PIL import Image
import numpy as np
from typing import Dict, List, Optional, Tuple
import albumentations as A
from collections import defaultdict
import cv2
from pycocotools import mask as mask_utils

# Constants
IMAGE_SIZE = 384  # Increased size for better detail capture
PATCH_SIZE = 14
BATCH_SIZE = 2
LEARNING_RATE = 1e-4
NUM_EPOCHS = 100
WARMUP_EPOCHS = 5


class FashionpediaDataset(Dataset):
    def __init__(self, data_dir: str, ann_file: str, transform=None, split='train'):
        self.data_dir = data_dir
        self.transform = transform
        self.split = split
        
        # Load annotations
        with open(ann_file, 'r') as f:
            self.data = json.load(f)
            
        # Create image id to annotations mapping
        self.img_to_anns = defaultdict(list)
        for ann in self.data['annotations']:
            self.img_to_anns[ann['image_id']].append(ann)
            
        # Create category and attribute mappings
        self.cat_to_id = {cat['name']: cat['id'] for cat in self.data['categories']}
        self.attr_to_id = {attr['name']: attr['id'] for attr in self.data['attributes']}
        
        self.images = self.data['images']
        
        # Image preprocessing
        self.image_processor = AutoImageProcessor.from_pretrained('facebook/dinov2-base')
        
    def __len__(self):
        return len(self.images)
        
    def __getitem__(self, idx):
        img_info = self.images[idx]
        img_path = os.path.join(self.data_dir, img_info['file_name'])
        image = Image.open(img_path).convert('RGB')
        
        # Get annotations for this image
        anns = self.img_to_anns[img_info['id']]
        
        # Prepare instance masks and labels
        masks = []
        boxes = []
        category_ids = []
        attribute_ids = []
        
        for ann in anns:
            mask = self._polygons_to_mask(ann['segmentation'], img_info['height'], img_info['width'])
            masks.append(mask)
            boxes.append(ann['bbox'])
            category_ids.append(ann['category_id'])
            attribute_ids.append(ann['attribute_ids'])
        
        # Process image with DINOv2 processor
        if self.transform:
            transformed = self.transform(image=np.array(image), 
                                      masks=masks,
                                      bboxes=boxes,
                                      category_ids=category_ids)
            
            image = transformed['image']
            masks = transformed['masks']
            boxes = transformed['bboxes']
        
        # Convert to tensors
        image = torch.FloatTensor(image)
        if image.ndim == 3:
            image = image.permute(2, 0, 1)
        masks = torch.FloatTensor(np.array(masks))
        boxes = torch.FloatTensor(np.array(boxes))
        category_ids = torch.LongTensor(category_ids)
        attribute_ids = [torch.LongTensor(attrs) for attrs in attribute_ids]
        
        return {
            'image': image,
            'masks': masks,
            'boxes': boxes,
            'category_ids': category_ids,
            'attribute_ids': attribute_ids,
            'image_id': img_info['id']
        }
    def _polygons_to_mask(self, segmentation, height, width):
        mask = np.zeros((height, width), dtype=np.uint8)
        
        # Handle RLE format 
        if isinstance(segmentation, dict):
            mask = mask_utils.decode(segmentation)
        # Handle polygon format
        elif isinstance(segmentation, list):
            polygons = []
            for seg in segmentation:
                poly = np.array([float(x) for x in seg]).astype(np.int32).reshape(-1, 2)
                polygons.append(poly)
                
            cv2.fillPoly(mask, polygons, 1)
        
        return mask

class MLP(nn.Module):
    """Multi-layer Perceptron"""
    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
        super().__init__()
        self.num_layers = num_layers
        h = [hidden_dim] * (num_layers - 1)
        self.layers = nn.ModuleList(
            nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])
        )
        
    def forward(self, x):
        for i, layer in enumerate(self.layers):
            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
        return x

class ExpertModule(nn.Module):
    def __init__(self, dim, expansion_factor=4):
        super().__init__()
        hidden_dim = dim * expansion_factor
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, dim)
        )
        
    def forward(self, x):
        return self.net(x)

class MixtureOfExperts(nn.Module):
    def __init__(self, dim, num_experts=4, k=2):
        super().__init__()
        self.num_experts = num_experts
        self.k = k
        
        self.experts = nn.ModuleList([ExpertModule(dim) for _ in range(num_experts)])
        self.gate = nn.Linear(dim, num_experts)
        
    def forward(self, x):
        batch_size, seq_len, dim = x.shape
        
        # Get expert weights
        gate_logits = self.gate(x)  # [batch_size, seq_len, num_experts]
        weights = F.softmax(gate_logits, dim=-1)
        
        # Select top-k experts
        top_k_weights, top_k_indices = torch.topk(weights, self.k, dim=-1)  # [batch_size, seq_len, k]
        top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True)
        
        # Compute expert outputs
        expert_outputs = []
        for expert in self.experts:
            expert_outputs.append(expert(x))
        expert_outputs = torch.stack(expert_outputs, dim=1)  # [batch_size, num_experts, seq_len, dim]
        
        # Prepare indices for gathering
        batch_indices = torch.arange(batch_size, device=x.device)[:, None, None, None]
        batch_indices = batch_indices.expand(batch_size, seq_len, self.k, 1)
        seq_indices = torch.arange(seq_len, device=x.device)[None, :, None, None]
        seq_indices = seq_indices.expand(batch_size, seq_len, self.k, 1)
        expert_indices = top_k_indices.unsqueeze(-1)
        
        # Gather the selected expert outputs
        gather_indices = torch.cat([batch_indices, expert_indices, seq_indices, 
                                  torch.zeros_like(batch_indices)], dim=-1)
        selected_outputs = expert_outputs[
            gather_indices[..., 0],
            gather_indices[..., 1],
            gather_indices[..., 2]
        ]  # [batch_size, seq_len, k, dim]
        
        # Apply weights and sum
        output = torch.sum(top_k_weights.unsqueeze(-1) * selected_outputs, dim=2)
        
        return output

class EfficientAttention(nn.Module):
    def __init__(self, dim, num_heads=8):
        super().__init__()
        self.num_heads = num_heads
        self.scale = (dim // num_heads) ** -0.5
        
        self.to_qkv = nn.Linear(dim, dim * 3, bias=False)
        self.to_out = nn.Linear(dim, dim)
        
    def forward(self, x):
        b, n, d = x.shape
        qkv = self.to_qkv(x).chunk(3, dim=-1)
        q, k, v = map(lambda t: t.view(b, n, self.num_heads, -1).transpose(1, 2), qkv)
        
        # Efficient attention computation
        dots = torch.matmul(q, k.transpose(-2, -1)) * self.scale
        attn = dots.softmax(dim=-1)
        
        out = torch.matmul(attn, v)
        out = out.transpose(1, 2).reshape(b, n, -1)
        return self.to_out(out)

class HierarchicalVisualEncoder(nn.Module):
    def __init__(self, hidden_dim=768):
        super().__init__()
        # Initialize DINOv2 backbone and processor
        self.processor = AutoImageProcessor.from_pretrained('facebook/dinov2-base')
        self.backbone = AutoModel.from_pretrained('facebook/dinov2-base')
        
        # Additional hierarchical processing
        self.hierarchical_transformer = nn.ModuleList([
            nn.TransformerEncoderLayer(d_model=hidden_dim, nhead=16)
            for _ in range(3)
        ])
        
        self.moe = MixtureOfExperts(hidden_dim)
        self.efficient_attention = EfficientAttention(hidden_dim)
        
    def forward(self, x):
        # Denormalize the images from [-mean/std, 1-mean/std] back to [0, 1]
        mean = torch.tensor([0.485, 0.456, 0.406], device=x.device).view(1, 3, 1, 1)
        std = torch.tensor([0.229, 0.224, 0.225], device=x.device).view(1, 3, 1, 1)
        x = x * std + mean
        
        # Rescale to [0, 255]
        x = x * 255
        x = x.to(torch.uint8)
        
        # Process images using DINOv2 processor
        inputs = self.processor(images=x.cpu().numpy(), return_tensors="pt")
        inputs = {k: v.to(x.device) if isinstance(v, torch.Tensor) else v 
                 for k, v in inputs.items()}
        
        # Get DINOv2 features
        outputs = self.backbone(**inputs)
        features = outputs.last_hidden_state
        
        # Hierarchical processing
        for transformer_layer in self.hierarchical_transformer:
            features = transformer_layer(features)
        
        # Apply MoE and efficient attention
        features = self.moe(features)
        features = self.efficient_attention(features)
        
        return features

class AdvancedMaskDecoder(nn.Module):
    def __init__(self, hidden_dim, num_classes, num_queries=100):
        super().__init__()
        self.num_queries = num_queries
        self.hidden_dim = hidden_dim
        
        # Learnable queries
        self.query_embed = nn.Embedding(num_queries, hidden_dim)
        
        # Positional encoding with fixed size
        self.pos_encoding = nn.Parameter(torch.randn(1, num_queries, hidden_dim))
        
        # Multi-scale decoder
        self.decoder_layers = nn.ModuleList([
            nn.TransformerDecoderLayer(
                d_model=hidden_dim,
                nhead=8,
                dim_feedforward=hidden_dim * 4,
                dropout=0.1,
                activation='gelu'
            )
            for _ in range(6)
        ])
        
        # Output heads
        self.mask_head = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, IMAGE_SIZE * IMAGE_SIZE)  # Changed to output full-size masks
        )
        
        self.box_head = MLP(hidden_dim, hidden_dim, 4, 3)
        
    def forward(self, features, mask=None):
        batch_size = features.shape[0]
        
        # Create queries
        queries = self.query_embed.weight.unsqueeze(0).repeat(batch_size, 1, 1)
        queries = queries + self.pos_encoding
        
        # Process through decoder layers
        queries = queries.transpose(0, 1)
        features = features.transpose(0, 1)
        
        output = queries
        for layer in self.decoder_layers:
            if mask is not None:
                mask = mask.unsqueeze(1).expand(
                    batch_size,
                    layer.nhead,
                    self.num_queries,
                    features.size(0)
                )
            output = layer(output, features, memory_mask=mask)
        
        output = output.transpose(0, 1)
        
        # Generate predictions
        # Reshape mask predictions to match target dimensions
        masks = self.mask_head(output)  # [batch_size, num_queries, IMAGE_SIZE * IMAGE_SIZE]
        masks = masks.view(batch_size, self.num_queries, IMAGE_SIZE, IMAGE_SIZE)  # Reshape to spatial dimensions
        
        boxes = self.box_head(output).sigmoid()  # [batch_size, num_queries, 4]
        
        return masks, boxes

class DiceLoss(nn.Module):
    def forward(self, inputs, targets):
        # Ensure inputs and targets have the same shape
        # inputs: [batch_size, num_queries, H, W]
        # targets: [batch_size, num_instances, H, W]
        
        # Flatten spatial dimensions
        inputs = inputs.flatten(2)  # [batch_size, num_queries, H*W]
        targets = targets.flatten(2)  # [batch_size, num_instances, H*W]
        
        # Calculate similarity matrix between all predictions and all targets
        inputs = inputs.sigmoid()
        
        # Compute pairwise IoU between predictions and targets
        intersection = torch.matmul(inputs, targets.transpose(-2, -1))  # [batch_size, num_queries, num_instances]
        inputs_sum = inputs.sum(-1).unsqueeze(-1)  # [batch_size, num_queries, 1]
        targets_sum = targets.sum(-1).unsqueeze(-2)  # [batch_size, 1, num_instances]
        union = inputs_sum + targets_sum - intersection
        
        # Calculate Dice coefficient
        dice = (2 * intersection + 1e-6) / (union + 1e-6)
        
        # Take the maximum dice coefficient for each target
        dice_max = dice.max(dim=1)[0]  # [batch_size, num_instances]
        
        # Return negative mean dice coefficient
        return 1 - dice_max.mean()

class AttributePredictor(nn.Module):
    def __init__(self, hidden_dim, num_attributes):
        super().__init__()
        self.num_attributes = num_attributes
        
        # Instance-level attribute prediction
        self.instance_attr = nn.Sequential(
            nn.LayerNorm(hidden_dim),  # Add normalization for better stability
            nn.Linear(hidden_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(0.1),  # Add dropout for regularization
            nn.Linear(hidden_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, num_attributes)
        )
        
        # Global attribute prediction with attention
        self.attention = nn.MultiheadAttention(hidden_dim, num_heads=8, batch_first=True)
        self.global_attr = nn.Sequential(
            nn.LayerNorm(hidden_dim),
            nn.Linear(hidden_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, num_attributes)
        )
        
    def forward(self, features, num_instances=None):
        """
        Args:
            features: Tensor of shape [B, N, D] where B is batch size, N is sequence length,
                     D is hidden dimension
            num_instances: Optional tensor of shape [B] containing number of valid instances
                         per batch item
        Returns:
            instance_attrs: Tensor of shape [B, max_instances, num_attributes]
            global_attrs: Tensor of shape [B, num_attributes]
        """
        B, N, D = features.shape
        max_instances = 10  # Default maximum number of instances
        
        if num_instances is not None:
            # Ensure num_instances is on the correct device
            num_instances = num_instances.to(features.device)
            # Get maximum number of instances in the batch
            max_instances = torch.max(num_instances).item()
        
        # Clip max_instances to sequence length
        max_instances = min(max_instances, N)
        
        # Process instance-level attributes
        instance_features = features[:, :max_instances, :]
        instance_attrs = self.instance_attr(instance_features)  # [B, max_instances, num_attributes]
        
        # Apply mask based on num_instances if provided
        if num_instances is not None:
            mask = torch.arange(max_instances, device=features.device)[None, :] >= num_instances[:, None]
            instance_attrs = instance_attrs.masked_fill(mask.unsqueeze(-1), 0.0)
        
        # Global attribute prediction with attention
        global_features, _ = self.attention(
            features,
            features,
            features,
            need_weights=False
        )
        
        # Pool global features with attention weights
        global_features = torch.mean(global_features, dim=1)  # [B, D]
        global_attrs = self.global_attr(global_features)  # [B, num_attributes]
        
        return instance_attrs, global_attrs

class AdvancedFashionTransformer(nn.Module):
    def __init__(self, num_classes, num_attributes, hidden_dim=1024):
        super().__init__()
        
        self.visual_encoder = HierarchicalVisualEncoder(hidden_dim)
        self.mask_decoder = AdvancedMaskDecoder(hidden_dim, num_classes)
        self.attribute_predictor = AttributePredictor(hidden_dim, num_attributes)
        
    def forward(self, images, num_instances=None):
        # Extract hierarchical features
        features = self.visual_encoder(images)
        
        # Generate masks and boxes
        masks, boxes = self.mask_decoder(features)
        
        # Predict attributes with num_instances
        instance_attrs, global_attrs = self.attribute_predictor(features, num_instances)
        
        return {
            'masks': masks,
            'boxes': boxes,
            'instance_attributes': instance_attrs,
            'global_attributes': global_attrs
        }

class FocalLoss(nn.Module):
    def __init__(self, alpha=0.25, gamma=2):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma
        
    def forward(self, inputs, targets):
        """
        Memory-efficient implementation of Focal Loss
        Args:
            inputs: Tensor of shape [B, num_queries, H, W]
            targets: Tensor of shape [B, num_instances]
        """
        B, Q, H, W = inputs.shape
        total_loss = 0
        valid_targets = 0
        
        # Process each batch item separately to save memory
        for b in range(B):
            # Get predictions and targets for current batch
            batch_inputs = inputs[b]  # [num_queries, H, W]
            batch_targets = targets[b]  # [num_instances]
            
            # Filter out padding (-1) from targets
            valid_mask = batch_targets >= 0
            batch_targets = batch_targets[valid_mask]
            
            if len(batch_targets) == 0:
                continue
                
            # Calculate loss for each valid target
            batch_loss = 0
            for target_idx in batch_targets:
                # Create target mask for current category
                target_mask = torch.zeros((H, W), device=inputs.device)
                target_mask = target_mask.view(-1)
                
                # Calculate BCE loss for each query
                query_losses = []
                for q in range(Q):
                    pred = batch_inputs[q].view(-1)
                    bce_loss = F.binary_cross_entropy_with_logits(
                        pred, 
                        target_mask,
                        reduction='none'
                    )
                    
                    # Apply focal loss formula
                    pt = torch.exp(-bce_loss)
                    focal = self.alpha * (1-pt)**self.gamma * bce_loss
                    query_losses.append(focal.mean())
                
                # Take minimum loss across queries for this target
                batch_loss += torch.stack(query_losses).min()
            
            if len(batch_targets) > 0:
                total_loss += batch_loss / len(batch_targets)
                valid_targets += 1
        
        # Average loss across valid batches
        return total_loss / max(valid_targets, 1)

import torch
from scipy.optimize import linear_sum_assignment

def box_iou(boxes1, boxes2):
    """
    Compute IoU between two sets of boxes
    boxes1, boxes2: [N, 4], [M, 4] where each box is [x1, y1, x2, y2]
    """
    area1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1])
    area2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1])
    
    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2]
    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]
    
    wh = (rb - lt).clamp(min=0)  # [N,M,2]
    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]
    
    union = area1[:, None] + area2 - inter
    
    iou = inter / (union + 1e-6)
    return iou

def hungarian_matching(pred_boxes, target_boxes, pred_logits=None, target_labels=None):
    """
    Assign predicted boxes to target boxes using Hungarian algorithm
    pred_boxes: [num_queries, 4]
    target_boxes: [num_targets, 4]
    """
    with torch.no_grad():
        # Calculate cost matrix based on L1 loss
        cost_bbox = torch.cdist(pred_boxes, target_boxes, p=1)
        
        # Calculate IoU cost
        iou = box_iou(pred_boxes, target_boxes)
        cost_giou = -iou
        
        # Combine costs
        C = cost_bbox + cost_giou
        
        # Run Hungarian algorithm
        C = C.cpu().numpy()
        row_ind, col_ind = linear_sum_assignment(C)  # Correctly unpack the two arrays
        
        # Convert to tensors and return as a single tuple
        return (torch.as_tensor(row_ind, dtype=torch.int64),
                torch.as_tensor(col_ind, dtype=torch.int64))

def compute_box_loss(pred_boxes, target_boxes):
    """
    Compute box loss with proper matching
    pred_boxes: [batch_size, num_queries, 4]
    target_boxes: [batch_size, num_targets, 4]
    """
    total_loss = 0
    num_boxes = 0
    
    # Process each item in the batch
    for pred, target in zip(pred_boxes, target_boxes):
        # Remove padding from target boxes
        valid_mask = (target.sum(dim=1) > 0)
        valid_target = target[valid_mask]
        
        if len(valid_target) == 0:
            continue
            
        # Match predictions to targets
        idx_pred, idx_target = hungarian_matching(pred, valid_target)
        
        if len(idx_pred) > 0:
            # Compute L1 loss for matched pairs
            loss = F.l1_loss(
                pred[idx_pred], 
                valid_target[idx_target], 
                reduction='sum'
            )
            
            total_loss += loss
            num_boxes += len(idx_pred)
    
    if num_boxes > 0:
        total_loss = total_loss / num_boxes
    
    return total_loss


def train_one_epoch(model, dataloader, optimizer, scheduler, scaler, device):
    model.train()
    total_loss = 0
    
    for batch_idx, batch in enumerate(dataloader):
        try:
            # Move data to device
            images = batch['image'].to(device)
            masks = batch['masks'].to(device)
            boxes = batch['boxes'].to(device)
            category_ids = batch['category_ids'].to(device)
            attribute_ids = batch['attribute_ids'].to(device)
            
            # Get number of valid instances for each sample in batch
            num_valid_instances = (category_ids != -1).sum(dim=1)  # [B]
            
            optimizer.zero_grad()
            
            with autocast(device_type=device.type if device.type != 'mps' else 'cpu'):
                # Pass num_valid_instances to model
                outputs = model(images, num_valid_instances)
                
                # Rest of the training loop remains the same
                mask_loss = DiceLoss()(outputs['masks'], masks)
                box_loss = compute_box_loss(outputs['boxes'], boxes)
                category_loss = FocalLoss()(outputs['masks'], category_ids)
                
                # Handle instance attribute loss
                instance_attrs = outputs['instance_attributes']
                instance_attr_loss = 0
                valid_instances = num_valid_instances.sum().item()
                
                if valid_instances > 0:
                    for b in range(len(num_valid_instances)):
                        n = num_valid_instances[b].item()
                        if n > 0:
                            pred_attrs = instance_attrs[b, :n]
                            target_attrs = attribute_ids[b, :n]
                            
                            instance_attr_loss += F.binary_cross_entropy_with_logits(
                                pred_attrs,
                                target_attrs,
                                reduction='sum'
                            )
                    
                    instance_attr_loss = instance_attr_loss / valid_instances
                
                # Global attribute loss
                global_attr_targets = attribute_ids.max(dim=1)[0]
                global_attr_loss = F.binary_cross_entropy_with_logits(
                    outputs['global_attributes'],
                    global_attr_targets.float()
                )
                
                loss = (mask_loss + box_loss + category_loss + 
                       0.5 * instance_attr_loss + 0.5 * global_attr_loss)
            
            # Rest of the training loop remains the same
            scaler.scale(loss).backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            scaler.step(optimizer)
            scaler.update()
            
            if scheduler is not None:
                scheduler.step()
                
            total_loss += loss.item()
            
            if batch_idx % 10 == 0 and torch.cuda.is_available():
                print(f"Batch {batch_idx}: Loss = {loss.item():.4f}")
                print(f"GPU Memory allocated: {torch.cuda.memory_allocated()/1024**3:.2f} GB")
                
        except RuntimeError as e:
            if "out of memory" in str(e):
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                print(f"WARNING: out of memory on batch {batch_idx}. Skipping batch.")
                continue
            else:
                raise e
    
    return total_loss / len(dataloader)

# Add gradient checkpointing to the model
def enable_gradient_checkpointing(model):
    """Enable gradient checkpointing for supported modules"""
    for module in model.modules():
        if hasattr(module, 'gradient_checkpointing'):
            module.gradient_checkpointing = True
def main():
    # Setup device
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # Create transforms - Remove normalization since DINOv2 will handle it
    train_transform = A.Compose([
        A.RandomResizedCrop(IMAGE_SIZE, IMAGE_SIZE),
        A.HorizontalFlip(p=0.5),
        A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
        # DINOv2 processor will handle normalization
        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ], bbox_params=A.BboxParams(format='coco', label_fields=['category_ids']))
    
    val_transform = A.Compose([
        A.Resize(IMAGE_SIZE, IMAGE_SIZE),
        # DINOv2 processor will handle normalization
        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ], bbox_params=A.BboxParams(format='coco', label_fields=['category_ids']))
    
    # Create datasets
    train_dataset = FashionpediaDataset(
        data_dir='/content/fashionpedia/train/',
        ann_file='/content/fashionpedia/instances_attributes_train2020.json',
        transform=train_transform,
        split='train'
    )
    
    val_dataset = FashionpediaDataset(
        data_dir='/content/fashionpedia/test/',
        ann_file='/content/fashionpedia/instances_attributes_val2020.json',
        transform=val_transform,
        split='val'
    )
    
    # Create dataloaders
    train_loader = DataLoader(
        train_dataset,
        batch_size=BATCH_SIZE,
        shuffle=True,
        num_workers=2,
        pin_memory=True,
        collate_fn=collate_fn
    )
    
    val_loader = DataLoader(
        val_dataset,
        batch_size=BATCH_SIZE,
        shuffle=False,
        num_workers=2,
        pin_memory=True,
        collate_fn=collate_fn
    )
    
    # Initialize model
    model = AdvancedFashionTransformer(
        num_classes=len(train_dataset.cat_to_id),
        num_attributes=len(train_dataset.attr_to_id),
        hidden_dim=768
    ).to(device)
    
    # Initialize optimizer and scheduler
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=LEARNING_RATE,
        weight_decay=0.01,
        betas=(0.9, 0.999)
    )
    
    # Cosine scheduler with warmup
    total_steps = NUM_EPOCHS * len(train_loader)
    warmup_steps = WARMUP_EPOCHS * len(train_loader)
    
    scheduler = get_cosine_schedule_with_warmup(
        optimizer,
        num_warmup_steps=warmup_steps,
        num_training_steps=total_steps
    )
    
    # Initialize gradient scaler for mixed precision training
    scaler = GradScaler('cuda')
    
    # Training loop
    best_val_loss = float('inf')
    for epoch in range(NUM_EPOCHS):
        print(f"Epoch {epoch+1}/{NUM_EPOCHS}")
        
        # Train
        train_loss = train_one_epoch(
            model=model,
            dataloader=train_loader,
            optimizer=optimizer,
            scheduler=scheduler,
            scaler=scaler,
            device=device
        )
        
        # Validate
        val_loss = validate(
            model=model,
            dataloader=val_loader,
            device=device
        )
        
        # Save checkpoint if validation loss improved
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            save_checkpoint(
                model=model,
                optimizer=optimizer,
                scheduler=scheduler,
                epoch=epoch,
                loss=val_loss,
                filename='best_model.pth'
            )
        
        print(f"Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}")

def validate(model, dataloader, device):
    model.eval()
    total_loss = 0
    
    with torch.no_grad():
        for batch in dataloader:
            images = batch['image'].to(device)
            masks = batch['masks'].to(device)
            boxes = batch['boxes'].to(device)
            category_ids = batch['category_ids'].to(device)
            attribute_ids = batch['attribute_ids'].to(device)
            
            num_valid_instances = (category_ids != -1).sum(dim=1)
            
            outputs = model(images)
            
            # Calculate validation losses
            mask_loss = DiceLoss()(outputs['masks'], masks)
            box_loss = compute_box_loss(outputs['boxes'], boxes)
            category_loss = FocalLoss()(outputs['masks'], category_ids)
            
            # Instance attribute loss
            instance_attrs = outputs['instance_attributes']
            instance_attr_loss = 0
            valid_instances = 0
            
            for b in range(len(num_valid_instances)):
                n = num_valid_instances[b].item()
                if n > 0:
                    pred_attrs = instance_attrs[b, :n]
                    target_attrs = attribute_ids[b, :n]
                    
                    instance_attr_loss += F.binary_cross_entropy_with_logits(
                        pred_attrs,
                        target_attrs,
                        reduction='sum'
                    )
                    valid_instances += n
            
            if valid_instances > 0:
                instance_attr_loss = instance_attr_loss / valid_instances
            
            # Global attribute loss
            global_attr_targets = attribute_ids.max(dim=1)[0]
            global_attr_loss = F.binary_cross_entropy_with_logits(
                outputs['global_attributes'],
                global_attr_targets.float()
            )
            
            loss = (mask_loss + box_loss + category_loss + 
                   0.5 * instance_attr_loss + 0.5 * global_attr_loss)
            
            total_loss += loss.item()
    
    return total_loss / len(dataloader)



def collate_fn(batch):
    """Custom collate function to handle variable-sized masks, boxes, and attributes"""
    # Initialize lists for batch items
    images = []
    masks = []
    boxes = []
    category_ids = []
    attribute_ids = []
    image_ids = []
    
    # Find maximum number of instances across all batch items
    max_instances = max(max(item['masks'].size(0), item['boxes'].size(0)) for item in batch)
    
    # Find maximum number of attributes for padding
    max_attrs = max(attr.size(0) for item in batch for attr in item['attribute_ids'])
    
    for item in batch:
        # Add image and image_id
        images.append(item['image'])
        image_ids.append(item['image_id'])
        
        # Get current sizes
        num_masks = item['masks'].size(0)
        num_boxes = item['boxes'].size(0)
        
        # Pad masks
        if num_masks < max_instances:
            mask_padding = torch.zeros(
                (max_instances - num_masks, *item['masks'].shape[1:]),
                dtype=item['masks'].dtype,
                device=item['masks'].device
            )
            padded_masks = torch.cat([item['masks'], mask_padding], dim=0)
        else:
            padded_masks = item['masks'][:max_instances]  # Truncate if larger
        masks.append(padded_masks)
        
        # Pad boxes
        if num_boxes < max_instances:
            box_padding = torch.zeros(
                (max_instances - num_boxes, 4),
                dtype=item['boxes'].dtype,
                device=item['boxes'].device
            )
            padded_boxes = torch.cat([item['boxes'], box_padding], dim=0)
        else:
            padded_boxes = item['boxes'][:max_instances]  # Truncate if larger
        boxes.append(padded_boxes)
        
        # Pad category_ids
        num_categories = len(item['category_ids'])
        if num_categories < max_instances:
            category_padding = torch.full(
                (max_instances - num_categories,),
                -1,  # Use -1 as padding value for categories
                dtype=item['category_ids'].dtype,
                device=item['category_ids'].device
            )
            padded_category_ids = torch.cat([item['category_ids'], category_padding], dim=0)
        else:
            padded_category_ids = item['category_ids'][:max_instances]  # Truncate if larger
        category_ids.append(padded_category_ids)
        
        # Handle attribute_ids - create padded tensor for each instance
        instance_attrs = []
        for attr_tensor in item['attribute_ids']:
            # Pad each attribute tensor to max_attrs
            if attr_tensor.size(0) < max_attrs:
                attr_padding = torch.zeros(max_attrs - attr_tensor.size(0), 
                                        dtype=attr_tensor.dtype,
                                        device=attr_tensor.device)
                padded_attr = torch.cat([attr_tensor, attr_padding])
            else:
                padded_attr = attr_tensor[:max_attrs]
            instance_attrs.append(padded_attr)
            
        # Pad number of instances
        while len(instance_attrs) < max_instances:
            instance_attrs.append(torch.zeros(max_attrs, 
                                           dtype=item['attribute_ids'][0].dtype,
                                           device=item['attribute_ids'][0].device))
            
        # Stack instance attributes into a single tensor
        stacked_attrs = torch.stack(instance_attrs)
        attribute_ids.append(stacked_attrs)
    
    # Create the batch dictionary with stacked tensors
    batch_dict = {
        'image': torch.stack(images),
        'boxes': torch.stack(boxes),
        'masks': torch.stack(masks),
        'category_ids': torch.stack(category_ids),
        'attribute_ids': torch.stack(attribute_ids),  # Now a tensor of shape [batch_size, max_instances, max_attrs]
        'image_id': image_ids
    }
    
    return batch_dict

def save_checkpoint(model, optimizer, scheduler, epoch, loss, filename):
    """Save model checkpoint"""
    checkpoint = {
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'scheduler_state_dict': scheduler.state_dict(),
        'loss': loss
    }
    torch.save(checkpoint, filename)

def get_cosine_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, num_cycles=0.5):
    """Create a schedule with a learning rate that decreases following the values of the cosine function"""
    def lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))

    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

if __name__ == '__main__':
    main()